From 151de63feb3181bcb89e30d57a880235dea629d5 Mon Sep 17 00:00:00 2001 From: Aaron Perez Date: Mon, 4 May 2026 23:58:29 -0500 Subject: [PATCH] feat(SKY-9463): open source Google Sheets connector auth (#5815) --- ..._app_origin_to_google_oauth_credentials.py | 27 + ...collapsed_consent_credential_in_one_row.py | 78 ++ skyvern/config.py | 26 + skyvern/forge/agent_functions.py | 65 +- skyvern/forge/api_app.py | 12 + skyvern/forge/forge_app.py | 27 + skyvern/forge/sdk/db/agent_db.py | 2 + skyvern/forge/sdk/db/id.py | 6 + skyvern/forge/sdk/db/models.py | 37 + .../forge/sdk/db/repositories/google_oauth.py | 315 ++++++ skyvern/forge/sdk/routes/google_oauth.py | 188 +++ skyvern/forge/sdk/routes/google_sheets.py | 270 +++++ skyvern/forge/sdk/schemas/google_oauth.py | 61 + skyvern/forge/sdk/schemas/google_sheets.py | 60 + .../sdk/services/google_oauth_service.py | 522 +++++++++ .../sdk/services/google_sheets_service.py | 548 +++++++++ tests/unit/force_stub_app.py | 1 + tests/unit/google/__init__.py | 0 .../test_agent_function_google_credentials.py | 159 +++ tests/unit/google/test_google_oauth.py | 1003 +++++++++++++++++ .../google/test_google_oauth_repository.py | 197 ++++ tests/unit/google/test_sheets.py | 318 ++++++ tests/unit/google/test_sheets_headers.py | 161 +++ tests/unit/google/test_sheets_retry.py | 228 ++++ .../unit/google/test_sheets_route_helpers.py | 60 + tests/unit/google/test_sheets_runtime.py | 95 ++ 26 files changed, 4460 insertions(+), 6 deletions(-) create mode 100644 alembic/versions/2026_05_05_0445-53a306b96ed7_add_consent_app_origin_to_google_oauth_credentials.py create mode 100644 alembic/versions/2026_05_05_0445-768c42c4968a_add_google_oauth_credentials_collapsed_consent_credential_in_one_row.py create mode 100644 skyvern/forge/sdk/db/repositories/google_oauth.py create mode 100644 skyvern/forge/sdk/routes/google_oauth.py create mode 100644 skyvern/forge/sdk/routes/google_sheets.py create mode 100644 skyvern/forge/sdk/schemas/google_oauth.py create mode 100644 skyvern/forge/sdk/schemas/google_sheets.py create mode 100644 skyvern/forge/sdk/services/google_oauth_service.py create mode 100644 skyvern/forge/sdk/services/google_sheets_service.py create mode 100644 tests/unit/google/__init__.py create mode 100644 tests/unit/google/test_agent_function_google_credentials.py create mode 100644 tests/unit/google/test_google_oauth.py create mode 100644 tests/unit/google/test_google_oauth_repository.py create mode 100644 tests/unit/google/test_sheets.py create mode 100644 tests/unit/google/test_sheets_headers.py create mode 100644 tests/unit/google/test_sheets_retry.py create mode 100644 tests/unit/google/test_sheets_route_helpers.py create mode 100644 tests/unit/google/test_sheets_runtime.py diff --git a/alembic/versions/2026_05_05_0445-53a306b96ed7_add_consent_app_origin_to_google_oauth_credentials.py b/alembic/versions/2026_05_05_0445-53a306b96ed7_add_consent_app_origin_to_google_oauth_credentials.py new file mode 100644 index 000000000..f21a7d2c9 --- /dev/null +++ b/alembic/versions/2026_05_05_0445-53a306b96ed7_add_consent_app_origin_to_google_oauth_credentials.py @@ -0,0 +1,27 @@ +"""add_consent_app_origin_to_google_oauth_credentials + +Revision ID: 53a306b96ed7 +Revises: 768c42c4968a +Create Date: 2026-05-05T04:45:56.788589+00:00 + +""" + +from typing import Sequence, Union + +import sqlalchemy as sa + +from alembic import op + +# revision identifiers, used by Alembic. +revision: str = "53a306b96ed7" +down_revision: Union[str, None] = "768c42c4968a" +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + op.add_column("google_oauth_credentials", sa.Column("consent_app_origin", sa.String(), nullable=True)) + + +def downgrade() -> None: + op.drop_column("google_oauth_credentials", "consent_app_origin") diff --git a/alembic/versions/2026_05_05_0445-768c42c4968a_add_google_oauth_credentials_collapsed_consent_credential_in_one_row.py b/alembic/versions/2026_05_05_0445-768c42c4968a_add_google_oauth_credentials_collapsed_consent_credential_in_one_row.py new file mode 100644 index 000000000..8150b8e72 --- /dev/null +++ b/alembic/versions/2026_05_05_0445-768c42c4968a_add_google_oauth_credentials_collapsed_consent_credential_in_one_row.py @@ -0,0 +1,78 @@ +"""add google_oauth_credentials (collapsed: consent + credential in one row) + +Revision ID: 768c42c4968a +Revises: b19685399f97 +Create Date: 2026-05-05T04:45:56.787945+00:00 + +""" + +from typing import Sequence, Union + +import sqlalchemy as sa + +from alembic import op + +# revision identifiers, used by Alembic. +revision: str = "768c42c4968a" +down_revision: Union[str, None] = "b19685399f97" +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + op.create_table( + "google_oauth_credentials", + sa.Column("id", sa.String(), nullable=False), + sa.Column("organization_id", sa.String(), nullable=False), + sa.Column("credential_name", sa.String(), nullable=False, server_default="Default"), + sa.Column("provider", sa.String(), nullable=False, server_default="google"), + sa.Column("state", sa.String(), nullable=False, server_default="pending_consent"), + sa.Column( + "scopes_requested", + sa.JSON(), + nullable=False, + server_default=sa.text("'[]'"), + ), + sa.Column( + "scopes_granted", + sa.JSON(), + nullable=False, + server_default=sa.text("'[]'"), + ), + sa.Column("encrypted_refresh_token", sa.String(), nullable=True), + sa.Column("encrypted_method", sa.String(), nullable=True), + sa.Column("consent_nonce", sa.String(), nullable=True), + sa.Column("consent_redirect_uri", sa.String(), nullable=True), + sa.Column("consent_code_verifier", sa.String(), nullable=True), + sa.Column("consent_expires_at", sa.DateTime(), nullable=True), + sa.Column("created_at", sa.DateTime(), nullable=False, server_default=sa.func.now()), + sa.Column("modified_at", sa.DateTime(), nullable=False, server_default=sa.func.now()), + sa.ForeignKeyConstraint(["organization_id"], ["organizations.organization_id"]), + sa.PrimaryKeyConstraint("id"), + sa.CheckConstraint( + "state IN ('pending_consent', 'active', 'revoked', 'error')", + name="ck_google_oauth_credentials_state", + ), + ) + op.create_index("ix_google_oauth_credentials_organization_id", "google_oauth_credentials", ["organization_id"]) + op.create_index( + "ix_google_oauth_credentials_state", + "google_oauth_credentials", + ["state"], + ) + # Partial unique index: a consent_nonce is only unique while it's in flight (pending_consent). + # Post-promotion we null it out, so active/revoked/error rows never collide here. + op.create_index( + "ux_google_oauth_credentials_consent_nonce", + "google_oauth_credentials", + ["consent_nonce"], + unique=True, + postgresql_where=sa.text("consent_nonce IS NOT NULL"), + ) + + +def downgrade() -> None: + op.drop_index("ux_google_oauth_credentials_consent_nonce", table_name="google_oauth_credentials") + op.drop_index("ix_google_oauth_credentials_state", table_name="google_oauth_credentials") + op.drop_index("ix_google_oauth_credentials_organization_id", table_name="google_oauth_credentials") + op.drop_table("google_oauth_credentials") diff --git a/skyvern/config.py b/skyvern/config.py index 73c294bb5..ef1304f01 100644 --- a/skyvern/config.py +++ b/skyvern/config.py @@ -524,6 +524,32 @@ class Settings(BaseSettings): ENCRYPTOR_AES_SECRET_KEY: str = "fillmein" ENCRYPTOR_AES_SALT: str | None = None ENCRYPTOR_AES_IV: str | None = None + ENABLE_ENCRYPTION: bool = False + + # Google OAuth settings (used by the Google Sheets connector) + GOOGLE_OAUTH_CLIENT_ID: str | None = None + GOOGLE_OAUTH_CLIENT_SECRET: str | None = None + # Hostnames allowed as the OAuth ``redirect_uri`` sent to Google. Defense-in-depth + # alongside Google's own redirect_uri allowlist (which is the real enforcement + # gate — it validates the full URI against its registered list). This setting is + # intentionally host-only: any path or port on an approved host passes. Empty + # list means no redirect_uri may be supplied at all (the route layer rejects + # with 400) when CLIENT_ID is set. + GOOGLE_OAUTH_REDIRECT_HOSTS: list[str] = Field(default_factory=list) + # Origins allowed as the bounce-back destination after OAuth callback. + # Never sent to Google. Two entry shapes — port handling differs: + # - Exact-match: ``https://host:port`` matches that exact origin (port included). + # ``https://app.example.com`` does NOT match ``https://app.example.com:8443``. + # - Suffix wildcard: ``*.foo.com`` matches any HTTPS hostname ending in ``.foo.com`` + # regardless of port (so preview deploys on non-default ports work). Rejects + # bare-suffix spoofs like ``attacker-foo.com``. + # Fails closed: an empty list rejects every app_origin, so self-hosted operators + # who want to use the bounce-back flow must populate this with at least one entry. + GOOGLE_OAUTH_APP_ORIGINS: list[str] = Field(default_factory=list) + + # Google Sheets API runtime tuning + GOOGLE_SHEETS_API_TIMEOUT_SECONDS: float = 30.0 + GOOGLE_SHEETS_API_MAX_RETRIES: int = 3 # Cleanup Cron Settings ENABLE_CLEANUP_CRON: bool = False diff --git a/skyvern/forge/agent_functions.py b/skyvern/forge/agent_functions.py index e244d263e..90918989d 100644 --- a/skyvern/forge/agent_functions.py +++ b/skyvern/forge/agent_functions.py @@ -20,6 +20,7 @@ from skyvern.forge.sdk.db.agent_db import AgentDB from skyvern.forge.sdk.models import Step, StepStatus from skyvern.forge.sdk.schemas.organizations import Organization from skyvern.forge.sdk.schemas.tasks import Task, TaskStatus +from skyvern.forge.sdk.services import google_oauth_service from skyvern.forge.sdk.trace import traced from skyvern.forge.sdk.workflow.models.block import BlockTypeVar from skyvern.webeye.actions.actions import Action @@ -707,12 +708,32 @@ class AgentFunction: organization_id: str, credential_id: str, ) -> str | None: - """Get a Google Sheets access token for the given credential. + """Mint a Google Sheets access token from the stored refresh token. - Returns None in OSS. Cloud override uses the OAuth service to - decrypt the stored refresh token and exchange it for an access token. + Returns None on any failure so callers can surface a reconnect prompt + instead of crashing. Cloud overrides this with an access-token cache + on top of the same backend. """ - return None + try: + secrets = await google_oauth_service.load_credential_secrets( + organization_id=organization_id, + credential_id=credential_id, + ) + return await google_oauth_service.access_token_from_secrets(secrets) + except google_oauth_service.EncryptionNotConfiguredError: + LOG.error( + "Google credential encryption is not configured; operators must enable ENABLE_ENCRYPTION", + organization_id=organization_id, + credential_id=credential_id, + ) + return None + except Exception: + LOG.exception( + "Failed to get Google Sheets credentials", + organization_id=organization_id, + credential_id=credential_id, + ) + return None async def get_google_workspace_credentials( self, @@ -720,8 +741,40 @@ class AgentFunction: credential_id: str, required_scopes: list[str] | None = None, ) -> object | None: - """OSS no-op; cloud override returns a refreshed google.oauth2.credentials.Credentials or None.""" - return None + """Return a refreshed ``google.oauth2.credentials.Credentials``, or None on failure. + + ``required_scopes`` is accepted for forward-compat with cloud overrides + that may enforce scope checks; the OSS path does not yet use it. + """ + if required_scopes: + # Debug-level so OSS operators don't see this on every call; scope + # enforcement lives on the cloud override and a future caller can + # grep for this message if they need to audit usage. + LOG.debug( + "required_scopes ignored by OSS get_google_workspace_credentials; cloud override gates this", + required_scopes=required_scopes, + credential_id=credential_id, + ) + try: + secrets = await google_oauth_service.load_credential_secrets( + organization_id=organization_id, + credential_id=credential_id, + ) + return await google_oauth_service.credentials_from_secrets(secrets) + except google_oauth_service.EncryptionNotConfiguredError: + LOG.error( + "Google credential encryption is not configured; operators must enable ENABLE_ENCRYPTION", + organization_id=organization_id, + credential_id=credential_id, + ) + return None + except Exception: + LOG.exception( + "Failed to get Google Workspace credentials", + organization_id=organization_id, + credential_id=credential_id, + ) + return None async def ensure_sheet_tab( self, diff --git a/skyvern/forge/api_app.py b/skyvern/forge/api_app.py index 24dfffab8..54f25c4ee 100644 --- a/skyvern/forge/api_app.py +++ b/skyvern/forge/api_app.py @@ -24,6 +24,8 @@ from skyvern.forge.sdk.core.skyvern_context import SkyvernContext from skyvern.forge.sdk.db.exceptions import NotFoundError from skyvern.forge.sdk.db.models import Base from skyvern.forge.sdk.routes import internal_auth +from skyvern.forge.sdk.routes.google_oauth import google_oauth_router +from skyvern.forge.sdk.routes.google_sheets import google_sheets_router from skyvern.forge.sdk.routes.routers import base_router, legacy_base_router, legacy_v2_router from skyvern.forge.sdk.services.local_org_auth_token_service import ( ensure_local_api_key, @@ -245,6 +247,16 @@ def create_api_app() -> FastAPI: fastapi_app.include_router(base_router, prefix="/v1") fastapi_app.include_router(legacy_base_router, prefix="/api/v1") fastapi_app.include_router(legacy_v2_router, prefix="/api/v2") + # Google OAuth is mounted on both the canonical ``/v1`` (matches base_router) + # and the legacy ``/api/v1`` (matches legacy_base_router) — same dual-prefix + # pattern as internal_auth below. Both registrations are intentional; do not + # remove either as a "duplicate". ``include_in_schema=False`` keeps the OAuth + # endpoints out of the public OpenAPI/Swagger surface — they're consumed by + # the frontend, not by SDK users. + fastapi_app.include_router(google_oauth_router, prefix="/v1/google", include_in_schema=False) + fastapi_app.include_router(google_oauth_router, prefix="/api/v1/google", include_in_schema=False) + fastapi_app.include_router(google_sheets_router, prefix="/v1/google/sheets", include_in_schema=False) + fastapi_app.include_router(google_sheets_router, prefix="/api/v1/google/sheets", include_in_schema=False) # local dev endpoints if settings.ENV == "local": diff --git a/skyvern/forge/forge_app.py b/skyvern/forge/forge_app.py index 80c206f85..873172249 100644 --- a/skyvern/forge/forge_app.py +++ b/skyvern/forge/forge_app.py @@ -24,6 +24,8 @@ from skyvern.forge.sdk.cache.base import BaseCache from skyvern.forge.sdk.cache.factory import CacheFactory from skyvern.forge.sdk.core.rate_limiter import NoopRateLimiter, RateLimiter from skyvern.forge.sdk.db.agent_db import AgentDB +from skyvern.forge.sdk.encrypt import encryptor +from skyvern.forge.sdk.encrypt.aes import AES from skyvern.forge.sdk.experimentation.providers import BaseExperimentationProvider, NoOpExperimentationProvider from skyvern.forge.sdk.schemas.credentials import CredentialVaultType from skyvern.forge.sdk.schemas.organizations import AzureClientSecretCredential, Organization @@ -116,6 +118,31 @@ def create_forge_app() -> ForgeApp: app.STORAGE = StorageFactory.get_storage() app.CACHE = CacheFactory.get_cache() + if settings.ENABLE_ENCRYPTION: + # Fail closed: a deployment that opts into encryption with the placeholder + # key would "encrypt" secrets with a public default — strictly worse than + # ENABLE_ENCRYPTION=false because operators would believe the data was + # protected. Salt/IV defaulting to None falls back to the deterministic + # ``default_salt`` / ``default_iv`` in ``aes.py`` — also a public default, + # so refuse the same way. + if not settings.ENCRYPTOR_AES_SECRET_KEY or settings.ENCRYPTOR_AES_SECRET_KEY == "fillmein": + raise RuntimeError( + "ENABLE_ENCRYPTION=true requires ENCRYPTOR_AES_SECRET_KEY to be set to a real " + "secret; empty values and the default placeholder 'fillmein' both fail closed." + ) + if not settings.ENCRYPTOR_AES_SALT or not settings.ENCRYPTOR_AES_IV: + raise RuntimeError( + "ENABLE_ENCRYPTION=true requires both ENCRYPTOR_AES_SALT and ENCRYPTOR_AES_IV to " + "be configured; unset values fall back to public defaults." + ) + encryptor.add_encrypt_method( + AES( + secret_key=settings.ENCRYPTOR_AES_SECRET_KEY, + salt=settings.ENCRYPTOR_AES_SALT, + iv=settings.ENCRYPTOR_AES_IV, + ) + ) + app.ARTIFACT_MANAGER = ArtifactManager() app.BROWSER_MANAGER = RealBrowserManager() app.EXPERIMENTATION_PROVIDER = NoOpExperimentationProvider() diff --git a/skyvern/forge/sdk/db/agent_db.py b/skyvern/forge/sdk/db/agent_db.py index 4cf56d3b9..2a8775579 100644 --- a/skyvern/forge/sdk/db/agent_db.py +++ b/skyvern/forge/sdk/db/agent_db.py @@ -20,6 +20,7 @@ from skyvern.forge.sdk.db.repositories.browser_sessions import BrowserSessionsRe from skyvern.forge.sdk.db.repositories.credentials import CredentialRepository from skyvern.forge.sdk.db.repositories.debug import DebugRepository from skyvern.forge.sdk.db.repositories.folders import FoldersRepository +from skyvern.forge.sdk.db.repositories.google_oauth import GoogleOAuthRepository from skyvern.forge.sdk.db.repositories.observer import ObserverRepository from skyvern.forge.sdk.db.repositories.organizations import OrganizationsRepository from skyvern.forge.sdk.db.repositories.otp import OTPRepository @@ -131,6 +132,7 @@ class AgentDB(BaseAlchemyDB): self.organizations = OrganizationsRepository(self.Session, debug_enabled, self.is_retryable_error) self.scripts = ScriptsRepository(self.Session, debug_enabled, self.is_retryable_error) self.browser_sessions = BrowserSessionsRepository(self.Session, debug_enabled, self.is_retryable_error) + self.google_oauth = GoogleOAuthRepository(self.Session, debug_enabled, self.is_retryable_error) self.schedules = SchedulesRepository( self.Session, debug_enabled, diff --git a/skyvern/forge/sdk/db/id.py b/skyvern/forge/sdk/db/id.py index b0ad45345..cf4ad8be7 100644 --- a/skyvern/forge/sdk/db/id.py +++ b/skyvern/forge/sdk/db/id.py @@ -44,6 +44,7 @@ CREDENTIAL_PREFIX = "cred" DEBUG_SESSION_PREFIX = "ds" FOLDER_PREFIX = "fld" BROWSER_PROFILE_PREFIX = "bp" +GOOGLE_OAUTH_CREDENTIAL_PREFIX = "goac" ORGANIZATION_BITWARDEN_COLLECTION_PREFIX = "obc" TASK_V2_ID = "tsk_v2" THOUGHT_ID = "ot" @@ -246,6 +247,11 @@ def generate_folder_id() -> str: return f"{FOLDER_PREFIX}_{int_id}" +def generate_google_oauth_credential_id() -> str: + int_id = generate_id() + return f"{GOOGLE_OAUTH_CREDENTIAL_PREFIX}_{int_id}" + + def generate_organization_bitwarden_collection_id() -> str: int_id = generate_id() return f"{ORGANIZATION_BITWARDEN_COLLECTION_PREFIX}_{int_id}" diff --git a/skyvern/forge/sdk/db/models.py b/skyvern/forge/sdk/db/models.py index 1dc9f5b99..482d53e3b 100644 --- a/skyvern/forge/sdk/db/models.py +++ b/skyvern/forge/sdk/db/models.py @@ -37,6 +37,7 @@ from skyvern.forge.sdk.db.id import ( generate_credential_parameter_id, generate_debug_session_id, generate_folder_id, + generate_google_oauth_credential_id, generate_onepassword_credential_parameter_id, generate_org_id, generate_organization_auth_token_id, @@ -1349,3 +1350,39 @@ class ScriptBranchHitModel(Base): hit_count = Column(Integer, default=1, nullable=False) first_hit_at = Column(DateTime, default=datetime.datetime.utcnow, nullable=False) last_hit_at = Column(DateTime, default=datetime.datetime.utcnow, nullable=False) + + +class GoogleOAuthCredentialModel(Base): + """Single-row lifecycle: pending_consent -> active -> revoked/error""" + + __tablename__ = "google_oauth_credentials" + __table_args__ = ( + Index( + "ux_google_oauth_credentials_consent_nonce", + "consent_nonce", + unique=True, + postgresql_where=text("consent_nonce IS NOT NULL"), + ), + ) + + id = Column(String, primary_key=True, default=generate_google_oauth_credential_id) + organization_id = Column(String, ForeignKey("organizations.organization_id"), index=True, nullable=False) + credential_name = Column(String, nullable=False, default="Default") + provider = Column(String, nullable=False, default="google") + state = Column(String, nullable=False, default="pending_consent", index=True) + scopes_requested = Column(JSON, nullable=False, default=list) + scopes_granted = Column(JSON, nullable=False, default=list) + encrypted_refresh_token = Column(String, nullable=True) + encrypted_method = Column(String, nullable=True) + consent_nonce = Column(String, nullable=True) + consent_redirect_uri = Column(String, nullable=True) + consent_code_verifier = Column(String, nullable=True) + consent_app_origin = Column(String, nullable=True) + consent_expires_at = Column(DateTime, nullable=True) + created_at = Column(DateTime, default=datetime.datetime.utcnow, nullable=False) + modified_at = Column( + DateTime, + default=datetime.datetime.utcnow, + onupdate=datetime.datetime.utcnow, + nullable=False, + ) diff --git a/skyvern/forge/sdk/db/repositories/google_oauth.py b/skyvern/forge/sdk/db/repositories/google_oauth.py new file mode 100644 index 000000000..850417610 --- /dev/null +++ b/skyvern/forge/sdk/db/repositories/google_oauth.py @@ -0,0 +1,315 @@ +from __future__ import annotations + +import datetime +from dataclasses import dataclass + +import structlog +from sqlalchemy import select, update + +from skyvern.forge.sdk.db._error_handling import db_operation +from skyvern.forge.sdk.db.base_repository import BaseRepository +from skyvern.forge.sdk.db.models import GoogleOAuthCredentialModel +from skyvern.forge.sdk.encrypt.base import EncryptMethod +from skyvern.forge.sdk.schemas.google_oauth import GoogleOAuthCredentialBase + +LOG = structlog.get_logger() + +# Credential lifecycle states. Kept as plain strings so DB rows survive code rewrites; +# the CHECK constraint in the migration pins the valid set. The repository owns these +# because the DB schema defines them — service re-exports for callers that previously +# imported through the service module. +STATE_PENDING_CONSENT = "pending_consent" +STATE_ACTIVE = "active" +STATE_REVOKED = "revoked" +STATE_ERROR = "error" + + +class InvalidConsentNonceError(ValueError): + """Raised when the OAuth callback nonce is unknown, expired, or already consumed. + + Defined here (not in the service) because ``promote_pending_to_active`` is the + only place it's raised — keeping it next to the raiser avoids the + service<->repo circular import that an in-method import previously dodged. + """ + + +@dataclass(frozen=True) +class PendingConsentContext: + credential_id: str + consent_redirect_uri: str | None + consent_code_verifier: str | None + consent_app_origin: str | None = None + + +@dataclass(frozen=True) +class ActiveCredentialCiphertext: + encrypted_refresh_token: str + encrypted_method: EncryptMethod + scopes_granted: list[str] + + +@dataclass(frozen=True) +class RevocableCiphertext: + exists: bool + encrypted_refresh_token: str | None = None + encrypted_method: EncryptMethod | None = None + + +class GoogleOAuthRepository(BaseRepository): + """All DB access for Google OAuth credentials. Owns session lifecycle.""" + + @db_operation("insert_pending_credential") + async def insert_pending_credential( + self, + credential_id: str, + organization_id: str, + credential_name: str, + scopes_requested: list[str], + consent_nonce: str, + consent_redirect_uri: str, + consent_expires_at: datetime.datetime, + consent_code_verifier: str, + consent_app_origin: str | None = None, + ) -> GoogleOAuthCredentialBase: + async with self.Session() as session: + model = GoogleOAuthCredentialModel( + id=credential_id, + organization_id=organization_id, + credential_name=credential_name, + provider="google", + state=STATE_PENDING_CONSENT, + scopes_requested=scopes_requested, + scopes_granted=[], + consent_nonce=consent_nonce, + consent_redirect_uri=consent_redirect_uri, + consent_expires_at=consent_expires_at, + consent_app_origin=consent_app_origin, + consent_code_verifier=consent_code_verifier, + ) + session.add(model) + await session.flush() + result = GoogleOAuthCredentialBase.model_validate(model, from_attributes=True) + await session.commit() + return result + + @db_operation("load_pending_by_nonce") + async def load_pending_by_nonce( + self, + organization_id: str, + nonce: str, + now: datetime.datetime | None = None, + ) -> PendingConsentContext | None: + # Reject expired rows here so the route layer doesn't burn Google's one-time + # auth code on a doomed exchange — promote_pending_to_active would catch the + # stale nonce, but only after the code has already been consumed. + cutoff = now if now is not None else datetime.datetime.now(datetime.UTC).replace(tzinfo=None) + async with self.Session() as session: + stmt = select( + GoogleOAuthCredentialModel.id, + GoogleOAuthCredentialModel.consent_redirect_uri, + GoogleOAuthCredentialModel.consent_code_verifier, + GoogleOAuthCredentialModel.consent_app_origin, + ).where( + GoogleOAuthCredentialModel.consent_nonce == nonce, + GoogleOAuthCredentialModel.organization_id == organization_id, + GoogleOAuthCredentialModel.state == STATE_PENDING_CONSENT, + GoogleOAuthCredentialModel.consent_expires_at >= cutoff, + ) + row = (await session.execute(stmt)).one_or_none() + if row is None: + return None + return PendingConsentContext( + credential_id=row[0], + consent_redirect_uri=row[1], + consent_code_verifier=row[2], + consent_app_origin=row[3], + ) + + @db_operation("promote_pending_to_active") + async def promote_pending_to_active( + self, + organization_id: str, + nonce: str, + encrypted_refresh_token: str, + encrypted_method: EncryptMethod, + scopes_granted: list[str], + now: datetime.datetime, + ) -> GoogleOAuthCredentialBase: + async with self.Session() as session: + stmt = ( + update(GoogleOAuthCredentialModel) + .where( + GoogleOAuthCredentialModel.consent_nonce == nonce, + GoogleOAuthCredentialModel.organization_id == organization_id, + GoogleOAuthCredentialModel.state == STATE_PENDING_CONSENT, + GoogleOAuthCredentialModel.consent_expires_at >= now, + ) + .values( + state=STATE_ACTIVE, + encrypted_refresh_token=encrypted_refresh_token, + encrypted_method=encrypted_method.value, + scopes_granted=scopes_granted, + consent_nonce=None, + consent_redirect_uri=None, + consent_expires_at=None, + consent_code_verifier=None, + consent_app_origin=None, + modified_at=now, + ) + .returning(GoogleOAuthCredentialModel) + ) + promoted = (await session.execute(stmt)).scalar_one_or_none() + if promoted is None: + fallback = ( + await session.execute( + select(GoogleOAuthCredentialModel).where( + GoogleOAuthCredentialModel.consent_nonce == nonce, + GoogleOAuthCredentialModel.organization_id == organization_id, + ) + ) + ).scalar_one_or_none() + if fallback is None: + raise InvalidConsentNonceError("Unknown OAuth consent nonce") + if fallback.state == STATE_ERROR: + raise InvalidConsentNonceError("OAuth consent row is in error state") + if fallback.state != STATE_PENDING_CONSENT: + raise InvalidConsentNonceError("OAuth consent nonce already consumed") + raise InvalidConsentNonceError("OAuth consent nonce expired") + result = GoogleOAuthCredentialBase.model_validate(promoted, from_attributes=True) + await session.commit() + LOG.info( + "Promoted pending Google OAuth credential", + credential_id=result.id, + organization_id=organization_id, + ) + return result + + @db_operation("list_active_for_org") + async def list_active_for_org(self, organization_id: str) -> list[GoogleOAuthCredentialBase]: + async with self.Session() as session: + stmt = ( + select(GoogleOAuthCredentialModel) + .where( + GoogleOAuthCredentialModel.organization_id == organization_id, + GoogleOAuthCredentialModel.state == STATE_ACTIVE, + ) + .order_by(GoogleOAuthCredentialModel.created_at.desc()) + ) + rows = (await session.execute(stmt)).scalars().all() + return [GoogleOAuthCredentialBase.model_validate(r, from_attributes=True) for r in rows] + + @db_operation("load_active_ciphertext") + async def load_active_ciphertext( + self, + organization_id: str, + credential_id: str, + ) -> ActiveCredentialCiphertext | None: + async with self.Session() as session: + stmt = select( + GoogleOAuthCredentialModel.encrypted_refresh_token, + GoogleOAuthCredentialModel.encrypted_method, + GoogleOAuthCredentialModel.scopes_granted, + ).where( + GoogleOAuthCredentialModel.id == credential_id, + GoogleOAuthCredentialModel.organization_id == organization_id, + GoogleOAuthCredentialModel.state == STATE_ACTIVE, + ) + row = (await session.execute(stmt)).one_or_none() + if row is None: + return None + ciphertext, method, scopes = row + if not ciphertext or not method: + return None + return ActiveCredentialCiphertext( + encrypted_refresh_token=ciphertext, + encrypted_method=EncryptMethod(method), + scopes_granted=list(scopes or []), + ) + + @db_operation("load_ciphertext_for_revoke") + async def load_ciphertext_for_revoke( + self, + organization_id: str, + credential_id: str, + ) -> RevocableCiphertext: + async with self.Session() as session: + stmt = select( + GoogleOAuthCredentialModel.encrypted_refresh_token, + GoogleOAuthCredentialModel.encrypted_method, + ).where( + GoogleOAuthCredentialModel.id == credential_id, + GoogleOAuthCredentialModel.organization_id == organization_id, + GoogleOAuthCredentialModel.state != STATE_REVOKED, + ) + row = (await session.execute(stmt)).one_or_none() + if row is None: + return RevocableCiphertext(exists=False) + ciphertext, method = row + if not ciphertext or not method: + return RevocableCiphertext(exists=True) + return RevocableCiphertext( + exists=True, + encrypted_refresh_token=ciphertext, + encrypted_method=EncryptMethod(method), + ) + + @db_operation("rename_active") + async def rename_active( + self, + organization_id: str, + credential_id: str, + credential_name: str, + now: datetime.datetime, + ) -> GoogleOAuthCredentialBase | None: + async with self.Session() as session: + stmt = ( + update(GoogleOAuthCredentialModel) + .where( + GoogleOAuthCredentialModel.id == credential_id, + GoogleOAuthCredentialModel.organization_id == organization_id, + GoogleOAuthCredentialModel.state == STATE_ACTIVE, + ) + .values(credential_name=credential_name, modified_at=now) + .returning(GoogleOAuthCredentialModel) + ) + row = (await session.execute(stmt)).scalar_one_or_none() + if row is None: + # No-match path: skip commit (nothing to persist) and let the + # session context manager close cleanly. Postgres has no row + # lock to release because the WHERE filtered out every row. + return None + result = GoogleOAuthCredentialBase.model_validate(row, from_attributes=True) + await session.commit() + return result + + @db_operation("mark_revoked_and_scrub") + async def mark_revoked_and_scrub( + self, + organization_id: str, + credential_id: str, + now: datetime.datetime, + ) -> str | None: + async with self.Session() as session: + stmt = ( + update(GoogleOAuthCredentialModel) + .where( + GoogleOAuthCredentialModel.id == credential_id, + GoogleOAuthCredentialModel.organization_id == organization_id, + GoogleOAuthCredentialModel.state != STATE_REVOKED, + ) + .values( + state=STATE_REVOKED, + encrypted_refresh_token=None, + encrypted_method=None, + consent_nonce=None, + consent_redirect_uri=None, + consent_expires_at=None, + consent_code_verifier=None, + consent_app_origin=None, + modified_at=now, + ) + .returning(GoogleOAuthCredentialModel.id) + ) + revoked_id = (await session.execute(stmt)).scalar_one_or_none() + await session.commit() + return revoked_id diff --git a/skyvern/forge/sdk/routes/google_oauth.py b/skyvern/forge/sdk/routes/google_oauth.py new file mode 100644 index 000000000..971176824 --- /dev/null +++ b/skyvern/forge/sdk/routes/google_oauth.py @@ -0,0 +1,188 @@ +from typing import Annotated + +import httpx +import requests # google-auth-oauthlib's Flow.fetch_token uses requests under the hood; we catch its transport errors. +import structlog +from fastapi import APIRouter, Depends, HTTPException +from google.auth.exceptions import GoogleAuthError +from oauthlib.oauth2 import InvalidGrantError, OAuth2Error + +from skyvern.forge.sdk.schemas.google_oauth import ( + CreateGoogleOAuthAuthorizeRequest, + CreateGoogleOAuthCallbackRequest, + GoogleOAuthAuthorizeResponse, + GoogleOAuthCredentialListResponse, + GoogleOAuthCredentialResponse, + UpdateGoogleOAuthCredentialRequest, +) +from skyvern.forge.sdk.schemas.organizations import Organization +from skyvern.forge.sdk.services import google_oauth_service, org_auth_service +from skyvern.forge.sdk.services.google_oauth_service import InvalidAppOriginError + +LOG = structlog.get_logger() + +google_oauth_router = APIRouter() + + +def _require_scopes_from_token(token_data: dict) -> list[str]: + """Extract the granted scopes from Google's token response, failing closed on empty/missing""" + raw = token_data.get("scope") + if raw is None: + raise HTTPException( + status_code=400, + detail="Google did not return any granted scopes. Please re-authorize and grant all requested scopes.", + ) + if isinstance(raw, str): + parts = [p for p in raw.replace(",", " ").split() if p] + if not parts: + raise HTTPException( + status_code=400, + detail="Google returned an empty scope. Please re-authorize and grant all requested scopes.", + ) + return parts + parts = [p for p in (str(s).strip() for s in raw) if p] + if not parts: + raise HTTPException( + status_code=400, + detail="Google returned no granted scopes. Please re-authorize and grant all requested scopes.", + ) + return parts + + +@google_oauth_router.post("/oauth/authorize") +async def google_oauth_authorize( + request: CreateGoogleOAuthAuthorizeRequest, + current_org: Annotated[Organization, Depends(org_auth_service.get_current_org)], +) -> GoogleOAuthAuthorizeResponse: + """Kick off the Google OAuth 2.0 authorization flow.""" + try: + start = await google_oauth_service.start_authorization( + organization_id=current_org.organization_id, + redirect_uri=request.redirect_uri, + credential_name=request.credential_name, + app_origin=request.app_origin, + ) + except InvalidAppOriginError as exc: + raise HTTPException(status_code=400, detail=str(exc)) + except google_oauth_service.InvalidRedirectURIError as exc: + raise HTTPException(status_code=400, detail=str(exc)) + except google_oauth_service.EncryptionNotConfiguredError as exc: + raise HTTPException(status_code=503, detail=str(exc)) + except ValueError as exc: + raise HTTPException(status_code=503, detail=str(exc)) + + return GoogleOAuthAuthorizeResponse(authorize_url=start.authorize_url, state=start.state) + + +@google_oauth_router.post("/oauth/callback") +async def google_oauth_callback( + request: CreateGoogleOAuthCallbackRequest, + current_org: Annotated[Organization, Depends(org_auth_service.get_current_org)], +) -> GoogleOAuthCredentialResponse: + """Handle the Google OAuth 2.0 authorization callback.""" + context = await google_oauth_service.load_pending_consent_context( + organization_id=current_org.organization_id, + nonce=request.state, + ) + if context is None or not context.consent_redirect_uri: + raise HTTPException(status_code=400, detail="Unknown or consumed OAuth consent nonce") + if not context.consent_code_verifier: + raise HTTPException( + status_code=400, + detail="OAuth consent row is missing the PKCE verifier; restart the consent flow", + ) + + try: + token_data = await google_oauth_service.exchange_code_for_tokens( + code=request.code, + redirect_uri=context.consent_redirect_uri, + code_verifier=context.consent_code_verifier, + ) + except ValueError as exc: + LOG.exception("Google OAuth client credentials not configured") + raise HTTPException(status_code=503, detail=str(exc)) + except InvalidGrantError as exc: + LOG.warning("Google OAuth invalid_grant on code exchange", error=str(exc)) + raise HTTPException(status_code=400, detail="Invalid or expired authorization code") + except OAuth2Error: + # Don't echo the exception string — OAuth2Error messages can carry the + # short-lived auth code or token-endpoint URLs. The full trace is in LOG.exception. + LOG.exception("Google OAuth protocol error on code exchange") + raise HTTPException(status_code=502, detail="Google OAuth exchange failed") + except (httpx.HTTPError, requests.RequestException, GoogleAuthError): + LOG.exception("Transport failure exchanging Google OAuth code") + raise HTTPException(status_code=502, detail="Upstream Google token endpoint failed") + except Exception: + LOG.exception("Unexpected failure exchanging Google OAuth code for tokens") + raise HTTPException(status_code=500, detail="Failed to exchange authorization code") + + refresh_token = token_data.get("refresh_token") + if not refresh_token: + raise HTTPException( + status_code=400, + detail="No refresh token received. Ensure access_type=offline and prompt=consent in the OAuth flow.", + ) + scopes_granted = _require_scopes_from_token(token_data) + + try: + credential = await google_oauth_service.promote_pending_credential( + organization_id=current_org.organization_id, + nonce=request.state, + refresh_token=refresh_token, + scopes_granted=scopes_granted, + ) + except google_oauth_service.InvalidConsentNonceError as exc: + raise HTTPException(status_code=400, detail=str(exc)) + except google_oauth_service.EncryptionNotConfiguredError as exc: + raise HTTPException(status_code=503, detail=str(exc)) + + # ``consent_app_origin`` was validated against ``GOOGLE_OAUTH_APP_ORIGINS`` in + # ``start_authorization`` before being persisted to the pending row, so reading + # it back here is safe. A future refactor that bypasses that pre-storage check + # must re-validate before echoing this value to the client. + return GoogleOAuthCredentialResponse(credential=credential, app_origin=context.consent_app_origin) + + +@google_oauth_router.get("/oauth/credentials") +async def list_google_oauth_credentials( + current_org: Annotated[Organization, Depends(org_auth_service.get_current_org)], +) -> GoogleOAuthCredentialListResponse: + """Fetch a list of Google OAuth credentials associated with an organization.""" + credentials = await google_oauth_service.get_credentials_for_org( + organization_id=current_org.organization_id, + ) + return GoogleOAuthCredentialListResponse(credentials=credentials) + + +@google_oauth_router.patch("/oauth/credentials/{credential_id}") +async def rename_google_oauth_credential( + credential_id: str, + request: UpdateGoogleOAuthCredentialRequest, + current_org: Annotated[Organization, Depends(org_auth_service.get_current_org)], +) -> GoogleOAuthCredentialResponse: + """Renames an existing Google OAuth credential for the specified organization""" + updated = await google_oauth_service.rename_credential( + organization_id=current_org.organization_id, + credential_id=credential_id, + credential_name=request.credential_name, + ) + if updated is None: + raise HTTPException(status_code=404, detail="Credential not found") + return GoogleOAuthCredentialResponse(credential=updated) + + +@google_oauth_router.delete( + "/oauth/credentials/{credential_id}", +) +async def delete_google_oauth_credential( + credential_id: str, + current_org: Annotated[Organization, Depends(org_auth_service.get_current_org)], +) -> dict[str, bool]: + """Deletes a specific Google OAuth credential associated with an organization""" + revoked = await google_oauth_service.revoke_credential( + organization_id=current_org.organization_id, + credential_id=credential_id, + ) + if not revoked: + raise HTTPException(status_code=404, detail="Credential not found") + return {"success": True} diff --git a/skyvern/forge/sdk/routes/google_sheets.py b/skyvern/forge/sdk/routes/google_sheets.py new file mode 100644 index 000000000..d9eed97f0 --- /dev/null +++ b/skyvern/forge/sdk/routes/google_sheets.py @@ -0,0 +1,270 @@ +from typing import Annotated + +import structlog +from fastapi import APIRouter, Depends, HTTPException, Query +from sqlalchemy.exc import SQLAlchemyError + +from skyvern.forge.sdk.schemas.google_sheets import ( + CreateGoogleSheetTabRequest, + CreateGoogleSheetTabResponse, + CreateGoogleSpreadsheetRequest, + CreateGoogleSpreadsheetResponse, + GetSheetDimensionsResponse, + GetSheetHeadersResponse, + GoogleSheetTab, + GoogleSpreadsheetSummary, + ListGoogleSheetTabsResponse, + ListGoogleSpreadsheetsResponse, + SheetHeader, +) +from skyvern.forge.sdk.schemas.organizations import Organization +from skyvern.forge.sdk.services import google_oauth_service, google_sheets_service, org_auth_service +from skyvern.schemas.google_sheets import column_index_to_letter + +LOG = structlog.get_logger() + +google_sheets_router = APIRouter() + + +async def _mint_access_token(organization_id: str, credential_id: str) -> str: + """Load credential secrets, then refresh without holding the session.""" + try: + secrets = await google_oauth_service.load_credential_secrets( + organization_id=organization_id, + credential_id=credential_id, + ) + except ValueError as exc: + raise HTTPException(status_code=404, detail=str(exc)) + except SQLAlchemyError: + # DB connectivity / timeout: the credential may be fine, retry server-side. + LOG.exception("Database error loading Google credential", credential_id=credential_id) + raise HTTPException(status_code=503, detail="Database unavailable") + except Exception as exc: + # Remaining failures (decrypt errors from key rotation / corrupted ciphertext) + # surface as reconnect_required so callers route the user back through consent. + LOG.exception("Failed to load Google credential secrets", credential_id=credential_id) + raise HTTPException( + status_code=409, + detail={ + "code": "reconnect_required", + "message": f"Could not load stored Google credential: {exc}", + }, + ) + try: + return await google_oauth_service.access_token_from_secrets(secrets) + except google_oauth_service.MissingAccessTokenError as exc: + # Google rejected the refresh token (invalid_grant after revoke/expiry); + # only a fresh consent round can recover, so route to the reconnect flow. + raise HTTPException( + status_code=409, + detail={ + "code": "reconnect_required", + "message": str(exc), + }, + ) + except Exception as exc: + LOG.exception("Failed to refresh Google access token", credential_id=credential_id) + raise HTTPException(status_code=502, detail=f"Failed to refresh access token: {exc}") + + +def _translate_api_error(exc: google_sheets_service.GoogleSheetsAPIError) -> HTTPException: + if exc.status == 403 and exc.code == "reconnect_required": + return HTTPException( + status_code=409, + detail={ + "code": "reconnect_required", + "missing_scope": True, + "message": exc.message, + }, + ) + if exc.status == 404: + return HTTPException(status_code=404, detail=exc.message) + if exc.status == 429: + return HTTPException(status_code=429, detail=exc.message) + if exc.status in (400, 401, 403): + return HTTPException(status_code=exc.status, detail=exc.message) + return HTTPException(status_code=502, detail=exc.message) + + +@google_sheets_router.get("/spreadsheets") +async def list_spreadsheets( + credential_id: Annotated[str, Query(description="Stored Google OAuth credential id")], + current_org: Annotated[Organization, Depends(org_auth_service.get_current_org)], + q: Annotated[str | None, Query(description="Optional search query on spreadsheet name")] = None, + page_token: Annotated[str | None, Query(description="Drive v3 pagination token")] = None, + page_size: Annotated[int, Query(ge=1, le=100)] = 25, +) -> ListGoogleSpreadsheetsResponse: + access_token = await _mint_access_token(current_org.organization_id, credential_id) + try: + paged = await google_sheets_service.list_spreadsheets( + access_token=access_token, + query=q, + page_token=page_token, + page_size=page_size, + ) + except google_sheets_service.GoogleSheetsAPIError as exc: + raise _translate_api_error(exc) + return ListGoogleSpreadsheetsResponse( + spreadsheets=[ + GoogleSpreadsheetSummary( + id=s.id, + name=s.name, + modified_time=s.modified_time, + web_view_link=s.web_view_link, + ) + for s in paged.spreadsheets + ], + next_page_token=paged.next_page_token, + ) + + +@google_sheets_router.get("/spreadsheets/{spreadsheet_id}") +async def get_spreadsheet( + spreadsheet_id: str, + credential_id: Annotated[str, Query(description="Stored Google OAuth credential id")], + current_org: Annotated[Organization, Depends(org_auth_service.get_current_org)], +) -> GoogleSpreadsheetSummary: + """Look up a spreadsheet's Drive metadata by id so the block picker can show the + human name after a workflow reload (the stored payload only carries the URL).""" + access_token = await _mint_access_token(current_org.organization_id, credential_id) + try: + summary = await google_sheets_service.get_spreadsheet_summary( + access_token=access_token, + spreadsheet_id=spreadsheet_id, + ) + except ValueError as exc: + raise HTTPException(status_code=400, detail=str(exc)) + except google_sheets_service.GoogleSheetsAPIError as exc: + raise _translate_api_error(exc) + return GoogleSpreadsheetSummary( + id=summary.id, + name=summary.name, + modified_time=summary.modified_time, + web_view_link=summary.web_view_link, + ) + + +@google_sheets_router.get("/spreadsheets/{spreadsheet_id}/tabs") +async def get_spreadsheet_tabs( + spreadsheet_id: str, + credential_id: Annotated[str, Query(description="Stored Google OAuth credential id")], + current_org: Annotated[Organization, Depends(org_auth_service.get_current_org)], +) -> ListGoogleSheetTabsResponse: + access_token = await _mint_access_token(current_org.organization_id, credential_id) + try: + tabs = await google_sheets_service.get_spreadsheet_tabs( + access_token=access_token, + spreadsheet_id=spreadsheet_id, + ) + except ValueError as exc: + raise HTTPException(status_code=400, detail=str(exc)) + except google_sheets_service.GoogleSheetsAPIError as exc: + raise _translate_api_error(exc) + return ListGoogleSheetTabsResponse( + tabs=[GoogleSheetTab(sheet_id=t.sheet_id, title=t.title, index=t.index) for t in tabs] + ) + + +@google_sheets_router.get("/spreadsheets/{spreadsheet_id}/headers") +async def get_spreadsheet_headers( + spreadsheet_id: str, + sheet_title: Annotated[str, Query(description="Tab title; query param so titles containing '/' work")], + credential_id: Annotated[str, Query(description="Stored Google OAuth credential id")], + current_org: Annotated[Organization, Depends(org_auth_service.get_current_org)], +) -> GetSheetHeadersResponse: + access_token = await _mint_access_token(current_org.organization_id, credential_id) + try: + rows = await google_sheets_service.get_sheet_headers( + access_token=access_token, + spreadsheet_id=spreadsheet_id, + sheet_title=sheet_title, + ) + except ValueError as exc: + raise HTTPException(status_code=400, detail=str(exc)) + except google_sheets_service.GoogleSheetsAPIError as exc: + raise _translate_api_error(exc) + return GetSheetHeadersResponse(headers=[SheetHeader(letter=h.letter, name=h.name) for h in rows]) + + +@google_sheets_router.get("/spreadsheets/{spreadsheet_id}/dimensions") +async def get_spreadsheet_dimensions( + spreadsheet_id: str, + sheet_title: Annotated[str, Query(description="Tab title; query param so titles containing '/' work")], + credential_id: Annotated[str, Query(description="Stored Google OAuth credential id")], + current_org: Annotated[Organization, Depends(org_auth_service.get_current_org)], +) -> GetSheetDimensionsResponse: + """Return the named tab's grid size + row-1 headers so the editor can preview the destination.""" + access_token = await _mint_access_token(current_org.organization_id, credential_id) + try: + grid = await google_sheets_service.get_sheet_grid_properties( + access_token=access_token, + spreadsheet_id=spreadsheet_id, + sheet_title=sheet_title, + ) + if grid is None: + # Don't echo `sheet_title` back: the dimensions endpoint is org-auth-gated + # but a future caller could render this as HTML, so keep it safe-by-default. + raise HTTPException(status_code=404, detail="Sheet tab not found") + rows = await google_sheets_service.get_sheet_headers( + access_token=access_token, + spreadsheet_id=spreadsheet_id, + sheet_title=sheet_title, + ) + except ValueError as exc: + raise HTTPException(status_code=400, detail=str(exc)) + except google_sheets_service.GoogleSheetsAPIError as exc: + raise _translate_api_error(exc) + last_index = max(grid.column_count - 1, 0) + return GetSheetDimensionsResponse( + sheet_id=grid.sheet_id, + title=grid.title, + column_count=grid.column_count, + row_count=grid.row_count, + last_column_letter=column_index_to_letter(last_index), + headers=[SheetHeader(letter=h.letter, name=h.name) for h in rows], + ) + + +@google_sheets_router.post("/spreadsheets") +async def create_spreadsheet( + request: CreateGoogleSpreadsheetRequest, + current_org: Annotated[Organization, Depends(org_auth_service.get_current_org)], +) -> CreateGoogleSpreadsheetResponse: + access_token = await _mint_access_token(current_org.organization_id, request.credential_id) + try: + created = await google_sheets_service.create_spreadsheet( + access_token=access_token, + title=request.title, + ) + except google_sheets_service.GoogleSheetsAPIError as exc: + raise _translate_api_error(exc) + return CreateGoogleSpreadsheetResponse( + spreadsheet=GoogleSpreadsheetSummary( + id=created.id, + name=created.title, + web_view_link=created.web_view_link, + ), + first_sheet_name=created.first_sheet_name, + ) + + +@google_sheets_router.post("/spreadsheets/{spreadsheet_id}/tabs") +async def create_sheet_tab( + spreadsheet_id: str, + request: CreateGoogleSheetTabRequest, + current_org: Annotated[Organization, Depends(org_auth_service.get_current_org)], +) -> CreateGoogleSheetTabResponse: + access_token = await _mint_access_token(current_org.organization_id, request.credential_id) + try: + tab = await google_sheets_service.create_sheet_tab( + access_token=access_token, + spreadsheet_id=spreadsheet_id, + title=request.title, + ) + except ValueError as exc: + raise HTTPException(status_code=400, detail=str(exc)) + except google_sheets_service.GoogleSheetsAPIError as exc: + raise _translate_api_error(exc) + return CreateGoogleSheetTabResponse( + tab=GoogleSheetTab(sheet_id=tab.sheet_id, title=tab.title, index=tab.index), + ) diff --git a/skyvern/forge/sdk/schemas/google_oauth.py b/skyvern/forge/sdk/schemas/google_oauth.py new file mode 100644 index 000000000..85a7b22b9 --- /dev/null +++ b/skyvern/forge/sdk/schemas/google_oauth.py @@ -0,0 +1,61 @@ +from datetime import datetime + +from pydantic import BaseModel, Field, field_validator + + +class GoogleOAuthCredentialBase(BaseModel): + id: str + organization_id: str + credential_name: str + provider: str = "google" + state: str + scopes_requested: list[str] = Field(default_factory=list) + scopes_granted: list[str] = Field(default_factory=list) + created_at: datetime + modified_at: datetime + + +class GoogleOAuthCredentialResponse(BaseModel): + credential: GoogleOAuthCredentialBase + app_origin: str | None = None + + +class GoogleOAuthCredentialListResponse(BaseModel): + credentials: list[GoogleOAuthCredentialBase] + + +class CreateGoogleOAuthAuthorizeRequest(BaseModel): + redirect_uri: str = Field(..., description="Redirect URI the consent flow will return to") + credential_name: str = Field(default="Default", description="Human-readable name for this credential") + app_origin: str | None = Field( + default=None, + description="Origin the callback should bounce back to (e.g. a Vercel preview URL). " + "Must be on the GOOGLE_OAUTH_APP_ORIGINS allowlist.", + ) + + +class GoogleOAuthAuthorizeResponse(BaseModel): + authorize_url: str + state: str + + +class CreateGoogleOAuthCallbackRequest(BaseModel): + code: str = Field(..., description="Authorization code from Google OAuth consent flow") + state: str = Field(..., description="State token returned by Google; must match a pending authorize request") + + +class UpdateGoogleOAuthCredentialRequest(BaseModel): + credential_name: str = Field( + ..., + min_length=1, + max_length=128, + description="New human-readable name for this credential", + ) + + @field_validator("credential_name") + @classmethod + def _strip_and_reject_blank(cls, v: str) -> str: + stripped = v.strip() + if not stripped: + raise ValueError("credential_name must not be blank") + return stripped diff --git a/skyvern/forge/sdk/schemas/google_sheets.py b/skyvern/forge/sdk/schemas/google_sheets.py new file mode 100644 index 000000000..612e751e9 --- /dev/null +++ b/skyvern/forge/sdk/schemas/google_sheets.py @@ -0,0 +1,60 @@ +from pydantic import BaseModel, Field + + +class GoogleSpreadsheetSummary(BaseModel): + id: str + name: str + modified_time: str | None = None + web_view_link: str | None = None + + +class ListGoogleSpreadsheetsResponse(BaseModel): + spreadsheets: list[GoogleSpreadsheetSummary] + next_page_token: str | None = None + + +class GoogleSheetTab(BaseModel): + sheet_id: int + title: str + index: int | None = None + + +class ListGoogleSheetTabsResponse(BaseModel): + tabs: list[GoogleSheetTab] + + +class SheetHeader(BaseModel): + letter: str + name: str + + +class GetSheetHeadersResponse(BaseModel): + headers: list[SheetHeader] + + +class GetSheetDimensionsResponse(BaseModel): + sheet_id: int + title: str + column_count: int + row_count: int + last_column_letter: str + headers: list[SheetHeader] + + +class CreateGoogleSpreadsheetRequest(BaseModel): + credential_id: str = Field(..., description="Stored Google OAuth credential id") + title: str = Field(..., min_length=1, max_length=255, description="Spreadsheet title") + + +class CreateGoogleSpreadsheetResponse(BaseModel): + spreadsheet: GoogleSpreadsheetSummary + first_sheet_name: str | None = None + + +class CreateGoogleSheetTabRequest(BaseModel): + credential_id: str = Field(..., description="Stored Google OAuth credential id") + title: str = Field(..., min_length=1, max_length=255, description="Tab title") + + +class CreateGoogleSheetTabResponse(BaseModel): + tab: GoogleSheetTab diff --git a/skyvern/forge/sdk/services/google_oauth_service.py b/skyvern/forge/sdk/services/google_oauth_service.py new file mode 100644 index 000000000..dd6fb09dd --- /dev/null +++ b/skyvern/forge/sdk/services/google_oauth_service.py @@ -0,0 +1,522 @@ +"""Open source Google OAuth 2.0 connector for Skyvern. + +This module implements the OAuth flow Skyvern uses to obtain a Google refresh +token on behalf of an organization, store it encrypted, and mint short-lived +access tokens for downstream API calls (Sheets, Drive, etc.). It is the +backend-agnostic implementation; the cloud deployment layers an in-process +access-token cache on top via ``CloudAgentFunction.get_google_sheets_credentials``. +""" + +from __future__ import annotations + +import asyncio +import datetime +import secrets +from dataclasses import dataclass, field +from urllib.parse import urlparse + +import httpx +import structlog +from google.auth.exceptions import GoogleAuthError +from google.auth.transport.requests import Request as GoogleAuthRequest +from google.oauth2.credentials import Credentials +from google_auth_oauthlib.flow import Flow + +from skyvern.config import settings +from skyvern.forge import app +from skyvern.forge.sdk.db.id import generate_google_oauth_credential_id + +# Lifecycle constants and InvalidConsentNonceError live on the repository (the DB owns +# them); re-export from this module so the cloud shim and existing call sites keep +# importing through ``google_oauth_service`` unchanged. +from skyvern.forge.sdk.db.repositories.google_oauth import ( # noqa: F401 + STATE_ACTIVE, + STATE_ERROR, + STATE_PENDING_CONSENT, + STATE_REVOKED, + InvalidConsentNonceError, + PendingConsentContext, +) +from skyvern.forge.sdk.encrypt import encryptor +from skyvern.forge.sdk.encrypt.base import EncryptMethod +from skyvern.forge.sdk.schemas.google_oauth import GoogleOAuthCredentialBase + +LOG = structlog.get_logger() + +GOOGLE_AUTHORIZE_ENDPOINT = "https://accounts.google.com/o/oauth2/v2/auth" +GOOGLE_TOKEN_ENDPOINT = "https://oauth2.googleapis.com/token" +GOOGLE_REVOKE_ENDPOINT = "https://oauth2.googleapis.com/revoke" + +# Spreadsheets + Drive scopes for the picker UX: +# - spreadsheets: read/write sheet cells and create new spreadsheets +# - drive.file: list/read spreadsheets this integration created or the user picked +# - drive.metadata.readonly: list file metadata (required for the picker search) +GOOGLE_SHEETS_SCOPES: tuple[str, ...] = ( + "https://www.googleapis.com/auth/spreadsheets", + "https://www.googleapis.com/auth/drive.file", + "https://www.googleapis.com/auth/drive.metadata.readonly", +) + +CONSENT_TTL_SECONDS = 600 + + +def google_access_token_cache_key(organization_id: str, credential_id: str) -> str: + return f"google:access_token:{organization_id}:{credential_id}" + + +class EncryptionNotConfiguredError(RuntimeError): + """Raised when Google credentials cannot be stored because AES encryption is disabled.""" + + +class MissingAccessTokenError(RuntimeError): + """Raised when Google returns a 2xx token response without an access_token.""" + + +class InvalidRedirectURIError(ValueError): + """Raised when the supplied OAuth redirect_uri host is not on the allowlist.""" + + +class InvalidAppOriginError(ValueError): + """Raised when the supplied app_origin is not on the GOOGLE_OAUTH_APP_ORIGINS allowlist.""" + + +@dataclass(frozen=True) +class GoogleCredentialSecrets: + """Decrypted credential material needed to mint an access token.""" + + refresh_token: str + scopes: list[str] = field(default_factory=list) + + +def _require_encryption() -> None: + """Google credentials must be encrypted at rest; refuse to write them in plaintext.""" + if not settings.ENABLE_ENCRYPTION: + raise EncryptionNotConfiguredError( + "Google OAuth credentials require AES encryption. Set ENABLE_ENCRYPTION=true and " + "configure ENCRYPTOR_AES_SECRET_KEY/SALT/IV on the deployment." + ) + + +def _require_client_credentials() -> tuple[str, str]: + client_id = settings.GOOGLE_OAUTH_CLIENT_ID + client_secret = settings.GOOGLE_OAUTH_CLIENT_SECRET + if not client_id or not client_secret: + raise ValueError("Google OAuth client credentials are not configured") + return client_id, client_secret + + +def _coerce_scopes(scopes: str | list[str] | tuple[str, ...] | None) -> list[str]: + """Accept space- or comma-delimited scope strings, iterables, or None and return a clean list.""" + if scopes is None: + return list(GOOGLE_SHEETS_SCOPES) + if isinstance(scopes, str): + parts = [p for p in scopes.replace(",", " ").split() if p] + else: + parts = [p for p in (str(s).strip() for s in scopes) if p] + return parts or list(GOOGLE_SHEETS_SCOPES) + + +_LOOPBACK_HOSTS = frozenset({"localhost", "127.0.0.1", "::1"}) + + +def _validate_redirect_uri(redirect_uri: str) -> None: + """Validate redirect_uri against the allowlist; defense-in-depth alongside Google's check. + + Beyond the host allowlist, we enforce scheme=https for non-loopback hosts so + an attacker who controls ``http://allowed-host.com`` can't satisfy a + hostname-only allowlist intended for ``https://allowed-host.com``. + """ + allowlist = settings.GOOGLE_OAUTH_REDIRECT_HOSTS + if not allowlist: + # Empty allowlist + no client_id = OAuth not configured at all → silent + # pass-through (the consent flow itself will fail later in + # ``_require_client_credentials``). Empty allowlist + client_id = misconfig + # we must catch up-front so we don't leak a redirect_uri to Google unchecked. + if settings.GOOGLE_OAUTH_CLIENT_ID: + raise InvalidRedirectURIError( + "GOOGLE_OAUTH_REDIRECT_HOSTS must be configured when GOOGLE_OAUTH_CLIENT_ID is set" + ) + return + try: + parsed = urlparse(redirect_uri) + scheme = parsed.scheme + # ``parsed.hostname`` is already lowercased; lowercase the allowlist too so an + # operator who configures ``MyApp.Example.Com`` doesn't reject every redirect. + host = parsed.hostname or "" + except Exception as exc: + raise InvalidRedirectURIError(f"Invalid redirect_uri: {exc}") + if host not in {h.strip().lower() for h in allowlist if h and h.strip()}: + raise InvalidRedirectURIError(f"redirect_uri host not allowed: {host}") + if scheme != "https" and host not in _LOOPBACK_HOSTS: + raise InvalidRedirectURIError(f"redirect_uri must use https for non-loopback host: {host}") + + +def _validate_app_origin(app_origin: str) -> None: + """Validate app_origin against the GOOGLE_OAUTH_APP_ORIGINS allowlist. + + Entries are either: + - Exact-match: full origin string like ``https://app.skyvern.com`` or ``http://localhost:5173``. + The port (if any) must match exactly — ``https://app.skyvern.com`` does NOT match + ``https://app.skyvern.com:8443``. + - Suffix wildcard: ``*.foo.com`` — matches any hostname ending with ``.foo.com`` over https only. + Wildcards intentionally ignore the port (preview deploys often run on non-default ports); + ``*.vercel.app`` accepts ``app.vercel.app:3000``. Rejects bare-suffix spoofs + (``attacker-foo.com``, ``foo.com.evil.com``). + + Fails closed: empty allowlist always raises. + """ + allowlist = settings.GOOGLE_OAUTH_APP_ORIGINS + if not allowlist: + raise InvalidAppOriginError("GOOGLE_OAUTH_APP_ORIGINS is not configured; app_origin is not accepted") + + try: + parsed = urlparse(app_origin) + scheme = parsed.scheme + netloc = parsed.netloc + host = parsed.hostname or "" + except Exception as exc: + raise InvalidAppOriginError(f"Invalid app_origin URL: {exc}") + + if not scheme or not netloc: + raise InvalidAppOriginError(f"app_origin must include scheme and hostname: {app_origin!r}") + + canonical = f"{scheme}://{netloc}" + + for entry in allowlist: + stripped = entry.strip() + if not stripped: + continue + if stripped.startswith("*."): + # Wildcard entry: only match over https + if scheme != "https": + continue + suffix = stripped[1:].lower() # e.g. ".vercel.app" + # Match against ``host`` (hostname only, no port) so wildcards like + # ``*.vercel.app`` accept ``myapp.vercel.app:3000``. Hostname must end + # with the suffix AND not be the suffix itself (no bare-suffix match). + if host.endswith(suffix) and host != suffix.lstrip("."): + return + else: + # Exact match + if canonical == stripped: + return + + raise InvalidAppOriginError(f"app_origin not allowed: {app_origin!r}") + + +def build_authorize_url( + redirect_uri: str, + state: str, + scopes: str | list[str] | tuple[str, ...] | None = None, + code_verifier: str | None = None, +) -> tuple[str, str]: + """Assemble the Google OAuth 2.0 consent URL the browser should navigate to. + + Self-validates ``redirect_uri`` against the host allowlist (defense-in-depth) + so direct callers cannot bypass the check. ``start_authorization`` also + validates up-front to avoid leaking a pending DB row on rejection. + + ``code_verifier`` may be supplied so callers that want the verifier + persisted to the DB before the URL exists can do so atomically; if None, + the Flow autogenerates one and we return it. + """ + # Require both halves up-front; otherwise consent completes but /callback fails every time + # because exchange_code_for_tokens needs the secret for the code -> token roundtrip. + client_id, client_secret = _require_client_credentials() + _validate_redirect_uri(redirect_uri) + + scope_list = _coerce_scopes(scopes) + flow = Flow.from_client_config( + { + "web": { + "client_id": client_id, + "client_secret": client_secret, + "auth_uri": GOOGLE_AUTHORIZE_ENDPOINT, + "token_uri": GOOGLE_TOKEN_ENDPOINT, + } + }, + scopes=scope_list, + redirect_uri=redirect_uri, + state=state, + autogenerate_code_verifier=code_verifier is None, + ) + if code_verifier is not None: + flow.code_verifier = code_verifier + url, _returned_state = flow.authorization_url( + access_type="offline", + prompt="consent", + include_granted_scopes="true", + ) + # Flow.code_verifier is None when autogenerate is False AND we didn't set + # it; that path is unreachable here, but raise (not assert) so a future + # refactor can't silently leak a null verifier into the DB under -O. + if not flow.code_verifier: + raise RuntimeError("Flow did not produce a PKCE code_verifier") + return url, flow.code_verifier + + +@dataclass(frozen=True) +class GoogleAuthorizationStart: + authorize_url: str + state: str + + +async def start_authorization( + organization_id: str, + redirect_uri: str, + credential_name: str = "Default", + scopes_requested: str | list[str] | tuple[str, ...] | None = None, + app_origin: str | None = None, +) -> GoogleAuthorizationStart: + """Insert a pending consent row, build the authorize URL, persist the PKCE verifier.""" + _require_encryption() + _validate_redirect_uri(redirect_uri) + if app_origin is not None: + _validate_app_origin(app_origin) + + credential_id = generate_google_oauth_credential_id() + nonce = secrets.token_urlsafe(32) + # Pre-generate the PKCE verifier so the pending row lands with the verifier + # populated in a single DB write — a crash mid-flow can no longer leave a + # pending row that the callback would see as missing the verifier. + code_verifier = secrets.token_urlsafe(64) + now = datetime.datetime.now(datetime.UTC).replace(tzinfo=None) + + await app.DATABASE.google_oauth.insert_pending_credential( + credential_id=credential_id, + organization_id=organization_id, + credential_name=credential_name, + scopes_requested=_coerce_scopes(scopes_requested), + consent_nonce=nonce, + consent_redirect_uri=redirect_uri, + consent_expires_at=now + datetime.timedelta(seconds=CONSENT_TTL_SECONDS), + consent_app_origin=app_origin, + consent_code_verifier=code_verifier, + ) + + authorize_url, _ = build_authorize_url( + redirect_uri=redirect_uri, + state=nonce, + scopes=scopes_requested, + code_verifier=code_verifier, + ) + + return GoogleAuthorizationStart(authorize_url=authorize_url, state=nonce) + + +async def promote_pending_credential( + organization_id: str, + nonce: str, + refresh_token: str, + scopes_granted: str | list[str] | tuple[str, ...] | None, +) -> GoogleOAuthCredentialBase: + """Encrypt the refresh token and promote the matching pending row to active.""" + _require_encryption() + encrypted_token = await encryptor.encrypt(refresh_token, EncryptMethod.AES) + now = datetime.datetime.now(datetime.UTC).replace(tzinfo=None) + return await app.DATABASE.google_oauth.promote_pending_to_active( + organization_id=organization_id, + nonce=nonce, + encrypted_refresh_token=encrypted_token, + encrypted_method=EncryptMethod.AES, + scopes_granted=_coerce_scopes(scopes_granted), + now=now, + ) + + +async def load_pending_consent_context(organization_id: str, nonce: str) -> PendingConsentContext | None: + """Fetch the redirect_uri + PKCE verifier the callback needs to replay to Google.""" + return await app.DATABASE.google_oauth.load_pending_by_nonce( + organization_id=organization_id, + nonce=nonce, + ) + + +async def exchange_code_for_tokens(code: str, redirect_uri: str, code_verifier: str | None) -> dict: + """Exchange an OAuth authorization code for access and refresh tokens.""" + client_id, client_secret = _require_client_credentials() + + # ``code_verifier`` is passed to ``fetch_token`` (where the PKCE verification + # actually happens); ``Flow.from_client_config`` ignores the kwarg. + flow = Flow.from_client_config( + { + "web": { + "client_id": client_id, + "client_secret": client_secret, + "auth_uri": GOOGLE_AUTHORIZE_ENDPOINT, + "token_uri": GOOGLE_TOKEN_ENDPOINT, + } + }, + scopes=None, + redirect_uri=redirect_uri, + ) + await asyncio.to_thread(flow.fetch_token, code=code, code_verifier=code_verifier) + creds = flow.credentials + # creds.scopes reflects what we *requested* on this Flow (None here, since we pass scopes=None at callback); + # Google's actually-granted scopes live on creds.granted_scopes / flow.oauth2session.token["scope"]. + granted = creds.granted_scopes or flow.oauth2session.token.get("scope") or "" + if isinstance(granted, (list, tuple)): + granted = " ".join(granted) + return { + "access_token": creds.token, + "refresh_token": creds.refresh_token, + "scope": granted, + "expiry": creds.expiry.isoformat() if creds.expiry else None, + } + + +async def refresh_access_token(refresh_token: str) -> dict: + """Exchange a refresh token for a new access token via google-auth.""" + client_id, client_secret = _require_client_credentials() + + creds = Credentials( + token=None, + refresh_token=refresh_token, + token_uri=GOOGLE_TOKEN_ENDPOINT, + client_id=client_id, + client_secret=client_secret, + ) + try: + await asyncio.to_thread(creds.refresh, GoogleAuthRequest()) + except GoogleAuthError as exc: + raise MissingAccessTokenError(f"Google token refresh failed: {exc}") from exc + return { + "access_token": creds.token, + "expiry": creds.expiry.isoformat() if creds.expiry else None, + } + + +async def load_credential_secrets( + organization_id: str, + credential_id: str, +) -> GoogleCredentialSecrets: + """Fetch + decrypt an active credential's refresh token.""" + payload = await app.DATABASE.google_oauth.load_active_ciphertext( + organization_id=organization_id, + credential_id=credential_id, + ) + if payload is None: + raise ValueError(f"No active Google OAuth credential found: {credential_id}") + refresh_token = await encryptor.decrypt(payload.encrypted_refresh_token, payload.encrypted_method) + return GoogleCredentialSecrets(refresh_token=refresh_token, scopes=payload.scopes_granted) + + +async def access_token_from_secrets(credential_secrets: GoogleCredentialSecrets) -> str: + """Exchange loaded secrets for an access token. Network-only; no DB.""" + if not credential_secrets.refresh_token: + raise ValueError("OAuth credential is missing refresh_token") + token_data = await refresh_access_token(credential_secrets.refresh_token) + access_token = token_data.get("access_token") + if not access_token: + raise MissingAccessTokenError("Google token response did not include access_token") + return access_token + + +async def credentials_from_secrets(credential_secrets: GoogleCredentialSecrets) -> Credentials: + """Build a refreshed ``Credentials`` from already-decrypted secrets. Network-only; no DB.""" + client_id, client_secret = _require_client_credentials() + creds = Credentials( + token=None, + refresh_token=credential_secrets.refresh_token, + token_uri=GOOGLE_TOKEN_ENDPOINT, + client_id=client_id, + client_secret=client_secret, + scopes=credential_secrets.scopes, + ) + try: + await asyncio.to_thread(creds.refresh, GoogleAuthRequest()) + except GoogleAuthError as exc: + raise MissingAccessTokenError(f"Google token refresh failed: {exc}") from exc + if not creds.token: + raise MissingAccessTokenError("Google token response did not include access_token") + return creds + + +async def get_credentials_for_org(organization_id: str) -> list[GoogleOAuthCredentialBase]: + """List all active Google OAuth credentials for an organization (metadata only).""" + return await app.DATABASE.google_oauth.list_active_for_org(organization_id=organization_id) + + +# google-auth 2.x has no first-class async revoke helper; keep the 10-line httpx POST. +async def _revoke_refresh_token_at_google(refresh_token: str) -> None: + """Best-effort call to Google's revoke endpoint.""" + try: + async with httpx.AsyncClient() as client: + response = await client.post( + GOOGLE_REVOKE_ENDPOINT, + data={"token": refresh_token}, + headers={"Content-Type": "application/x-www-form-urlencoded"}, + timeout=10, + ) + if response.status_code >= 400: + LOG.warning( + "Google token revocation returned non-success", + status_code=response.status_code, + body=response.text[:200], + ) + except Exception as exc: + LOG.warning("Failed to revoke Google refresh token upstream", error=str(exc)) + + +async def rename_credential( + organization_id: str, + credential_id: str, + credential_name: str, +) -> GoogleOAuthCredentialBase | None: + """Rename an active Google OAuth credential. Returns None for not-found / wrong-state / wrong-org.""" + now = datetime.datetime.now(datetime.UTC).replace(tzinfo=None) + return await app.DATABASE.google_oauth.rename_active( + organization_id=organization_id, + credential_id=credential_id, + credential_name=credential_name, + now=now, + ) + + +async def revoke_credential(organization_id: str, credential_id: str) -> bool: + """Revoke upstream at Google, mark local row revoked, invalidate access-token cache.""" + payload = await app.DATABASE.google_oauth.load_ciphertext_for_revoke( + organization_id=organization_id, + credential_id=credential_id, + ) + if not payload.exists: + return False + + refresh_token: str | None = None + if payload.encrypted_refresh_token and payload.encrypted_method: + try: + refresh_token = await encryptor.decrypt(payload.encrypted_refresh_token, payload.encrypted_method) + except Exception as exc: + LOG.warning( + "Failed to decrypt refresh token for upstream revocation", + credential_id=credential_id, + error=str(exc), + ) + if refresh_token: + await _revoke_refresh_token_at_google(refresh_token) + + now = datetime.datetime.now(datetime.UTC).replace(tzinfo=None) + revoked_id = await app.DATABASE.google_oauth.mark_revoked_and_scrub( + organization_id=organization_id, + credential_id=credential_id, + now=now, + ) + + if revoked_id is not None: + try: + # 30s tombstone outlives any in-flight token refresh that beat the revoke commit. + cache_key = google_access_token_cache_key(organization_id, credential_id) + await app.CACHE.set(cache_key, "", ex=datetime.timedelta(seconds=30)) + except Exception: + LOG.warning( + "Failed to invalidate access-token cache on revoke", + credential_id=credential_id, + exc_info=True, + ) + LOG.info( + "Revoked Google OAuth credential", + credential_id=credential_id, + organization_id=organization_id, + ) + return True + return False diff --git a/skyvern/forge/sdk/services/google_sheets_service.py b/skyvern/forge/sdk/services/google_sheets_service.py new file mode 100644 index 000000000..27f253814 --- /dev/null +++ b/skyvern/forge/sdk/services/google_sheets_service.py @@ -0,0 +1,548 @@ +import asyncio +from dataclasses import dataclass, field +from datetime import datetime, timezone +from email.utils import parsedate_to_datetime +from typing import Any +from urllib.parse import quote + +import httpx +import structlog + +from skyvern.config import settings +from skyvern.schemas.google_sheets import GoogleSheetsAPIError, extract_spreadsheet_id, quote_sheet_name # noqa: F401 + +LOG = structlog.get_logger() + +DRIVE_API_BASE = "https://www.googleapis.com/drive/v3" +SHEETS_API_BASE = "https://sheets.googleapis.com/v4/spreadsheets" +MIME_SPREADSHEET = "application/vnd.google-apps.spreadsheet" + +_RECONNECT_SCOPE_CODE = "reconnect_required" + +_RETRYABLE_STATUSES = {429, 500, 502, 503, 504} +_DEFAULT_BACKOFF_SECONDS = 1.0 + + +def _compute_backoff(attempt: int, retry_after: str | None) -> float: + if retry_after: + value = retry_after.strip() + try: + return max(0.0, float(value)) + except ValueError: + pass + # RFC 7231 also permits HTTP-date; without this branch we silently fall + # through to exponential backoff and hammer a rate-limited endpoint. + try: + target = parsedate_to_datetime(value) + except (TypeError, ValueError): + target = None + if target is not None: + if target.tzinfo is None: + target = target.replace(tzinfo=timezone.utc) + delta = (target - datetime.now(timezone.utc)).total_seconds() + return max(0.0, delta) + return _DEFAULT_BACKOFF_SECONDS * (2 ** (attempt - 1)) + + +async def _request_with_retry( + method: str, + url: str, + *, + params: dict[str, Any] | None = None, + json_body: dict[str, Any] | None = None, + headers: dict[str, str], + idempotent: bool = True, +) -> httpx.Response: + """Issue an HTTP request with retry/backoff. ``idempotent=False`` means + we never replay a mutation - not on transport/timeout, not on 5xx, not on + 429. Google does not guarantee a 429 (or 5xx) was rejected before the + write landed, so replaying risks duplicate rows the caller can't see. + """ + last_response: httpx.Response | None = None + max_attempts = max(1, settings.GOOGLE_SHEETS_API_MAX_RETRIES) + async with httpx.AsyncClient(timeout=settings.GOOGLE_SHEETS_API_TIMEOUT_SECONDS) as client: + for attempt in range(1, max_attempts + 1): + try: + response = await client.request(method, url, params=params, json=json_body, headers=headers) + except (httpx.TransportError, httpx.TimeoutException) as exc: + # Non-idempotent + transport failure: ambiguous (request may have landed), + # so surface immediately instead of replaying the mutation. + if not idempotent or attempt == max_attempts: + raise GoogleSheetsAPIError( + status=503, + code="upstream_unavailable", + message=f"Google Sheets transport failure: {exc}", + ) from exc + await asyncio.sleep(_compute_backoff(attempt, None)) + continue + if response.status_code not in _RETRYABLE_STATUSES: + return response + if not idempotent: + return response + last_response = response + if attempt == max_attempts: + break + await asyncio.sleep(_compute_backoff(attempt, response.headers.get("Retry-After"))) + if last_response is None: + raise RuntimeError("retry loop exited without a response") + return last_response + + +@dataclass(frozen=True) +class SpreadsheetSummary: + id: str + name: str + modified_time: str | None = None + web_view_link: str | None = None + + +@dataclass(frozen=True) +class PagedSpreadsheets: + spreadsheets: list[SpreadsheetSummary] = field(default_factory=list) + next_page_token: str | None = None + + +@dataclass(frozen=True) +class SheetTab: + sheet_id: int + title: str + index: int | None = None + + +@dataclass(frozen=True) +class CreatedSpreadsheet: + id: str + title: str + web_view_link: str | None = None + first_sheet_name: str | None = None + + +@dataclass(frozen=True) +class SheetHeader: + letter: str + name: str + + +@dataclass(frozen=True) +class SheetGridProperties: + sheet_id: int + title: str + column_count: int + row_count: int + + +def _escape_drive_q(value: str) -> str: + """Escape user input for a Drive v3 q= filter (backslash + single quote).""" + return value.replace("\\", "\\\\").replace("'", "\\'") + + +def _build_drive_q(search_query: str | None) -> str: + terms = [f"mimeType = '{MIME_SPREADSHEET}'", "trashed = false"] + if search_query: + terms.append(f"name contains '{_escape_drive_q(search_query)}'") + return " and ".join(terms) + + +def _raise_for_error(response: httpx.Response) -> None: + if response.is_success: + return + status = response.status_code + payload: dict[str, Any] = {} + try: + payload = response.json() or {} + except ValueError: + pass + err = payload.get("error") if isinstance(payload, dict) else None + message = "Google API error" + code: str | None = None + if isinstance(err, dict): + message = err.get("message") or message + details = err.get("errors") + if isinstance(details, list) and details: + reason = details[0].get("reason") if isinstance(details[0], dict) else None + if reason: + code = reason + # Only route to reconnect when Google reports a scope/permissions error on the grant; + # file-level 403s (e.g. `insufficientFilePermissions`) cannot be resolved by reconnecting + # so they must keep their original code and surface as a plain 403. + if status == 403 and code in {"insufficientPermissions", "insufficientScopes"}: + code = _RECONNECT_SCOPE_CODE + else: + message = response.text[:500] or message + raise GoogleSheetsAPIError(status=status, code=code, message=message) + + +def _auth_headers(access_token: str) -> dict[str, str]: + return {"Authorization": f"Bearer {access_token}", "Accept": "application/json"} + + +async def list_spreadsheets( + access_token: str, + query: str | None = None, + page_token: str | None = None, + page_size: int = 25, +) -> PagedSpreadsheets: + """List spreadsheets via Drive v3 files.list, filtered to Google Sheets MIME.""" + params: dict[str, Any] = { + "q": _build_drive_q(query), + "pageSize": max(1, min(100, page_size)), + "fields": "nextPageToken,files(id,name,modifiedTime,webViewLink)", + "orderBy": "modifiedTime desc", + "spaces": "drive", + # Include spreadsheets that live in shared drives; without these, picker + # results omit any sheet not owned by the user, which breaks orgs that + # standardize on shared drives. + "supportsAllDrives": "true", + "includeItemsFromAllDrives": "true", + } + if page_token: + params["pageToken"] = page_token + + response = await _request_with_retry( + "GET", + f"{DRIVE_API_BASE}/files", + params=params, + headers=_auth_headers(access_token), + ) + _raise_for_error(response) + payload = response.json() + items = payload.get("files") or [] + return PagedSpreadsheets( + spreadsheets=[ + SpreadsheetSummary( + id=item["id"], + name=item.get("name") or item["id"], + modified_time=item.get("modifiedTime"), + web_view_link=item.get("webViewLink"), + ) + for item in items + if item.get("id") + ], + next_page_token=payload.get("nextPageToken"), + ) + + +async def get_spreadsheet_tabs(access_token: str, spreadsheet_id: str) -> list[SheetTab]: + """Fetch the tab (sheet) metadata for a spreadsheet via Sheets v4.""" + spreadsheet_id = extract_spreadsheet_id(spreadsheet_id) + params = {"fields": "sheets(properties(sheetId,title,index))"} + response = await _request_with_retry( + "GET", + f"{SHEETS_API_BASE}/{spreadsheet_id}", + params=params, + headers=_auth_headers(access_token), + ) + _raise_for_error(response) + payload = response.json() + tabs: list[SheetTab] = [] + for sheet in payload.get("sheets") or []: + props = sheet.get("properties") or {} + sheet_id = props.get("sheetId") + title = props.get("title") + if sheet_id is None or title is None: + continue + tabs.append(SheetTab(sheet_id=int(sheet_id), title=title, index=props.get("index"))) + return tabs + + +async def get_spreadsheet_summary(access_token: str, spreadsheet_id: str) -> SpreadsheetSummary: + """Fetch a single spreadsheet's Drive metadata by id, for block rehydration on reload.""" + spreadsheet_id = extract_spreadsheet_id(spreadsheet_id) + params = { + "fields": "id,name,modifiedTime,webViewLink", + # Mirror list_spreadsheets so sheets in shared drives resolve. + "supportsAllDrives": "true", + "includeItemsFromAllDrives": "true", + } + response = await _request_with_retry( + "GET", + f"{DRIVE_API_BASE}/files/{quote(spreadsheet_id, safe='')}", + params=params, + headers=_auth_headers(access_token), + ) + _raise_for_error(response) + payload = response.json() + file_id = payload.get("id") + if not file_id: + raise GoogleSheetsAPIError( + status=500, + code="malformed_response", + message="Drive response missing file id", + ) + return SpreadsheetSummary( + id=file_id, + name=payload.get("name") or file_id, + modified_time=payload.get("modifiedTime"), + web_view_link=payload.get("webViewLink"), + ) + + +async def create_spreadsheet(access_token: str, title: str) -> CreatedSpreadsheet: + """Create a new spreadsheet via Sheets v4 spreadsheets.create.""" + body = {"properties": {"title": title or "Untitled spreadsheet"}} + response = await _request_with_retry( + "POST", + SHEETS_API_BASE, + json_body=body, + headers={**_auth_headers(access_token), "Content-Type": "application/json"}, + idempotent=False, + ) + _raise_for_error(response) + payload = response.json() + spreadsheet_id = payload.get("spreadsheetId") + if not spreadsheet_id: + raise GoogleSheetsAPIError( + status=500, + code="malformed_response", + message="Sheets API response missing spreadsheetId", + ) + props = payload.get("properties") or {} + sheets = payload.get("sheets") or [] + first_sheet_name: str | None = None + if sheets: + first_sheet_name = (sheets[0].get("properties") or {}).get("title") + return CreatedSpreadsheet( + id=spreadsheet_id, + title=props.get("title") or title, + web_view_link=payload.get("spreadsheetUrl"), + first_sheet_name=first_sheet_name, + ) + + +async def create_sheet_tab(access_token: str, spreadsheet_id: str, title: str) -> SheetTab: + """Append a new tab to a spreadsheet via batchUpdate addSheet.""" + spreadsheet_id = extract_spreadsheet_id(spreadsheet_id) + body = { + "requests": [ + { + "addSheet": { + "properties": {"title": title}, + } + } + ] + } + response = await _request_with_retry( + "POST", + f"{SHEETS_API_BASE}/{spreadsheet_id}:batchUpdate", + json_body=body, + headers={**_auth_headers(access_token), "Content-Type": "application/json"}, + idempotent=False, + ) + _raise_for_error(response) + payload = response.json() + replies = payload.get("replies") or [] + if not replies: + raise GoogleSheetsAPIError(status=500, code="malformed_response", message="batchUpdate returned no replies") + props = (replies[0].get("addSheet") or {}).get("properties") or {} + sheet_id = props.get("sheetId") + out_title = props.get("title") + if sheet_id is None or out_title is None: + raise GoogleSheetsAPIError(status=500, code="malformed_response", message="addSheet reply missing properties") + return SheetTab(sheet_id=int(sheet_id), title=out_title, index=props.get("index")) + + +def _column_index_to_letter(index: int) -> str: + # 0 -> "A", 25 -> "Z", 26 -> "AA" + letters: list[str] = [] + n = index + while True: + letters.append(chr(ord("A") + (n % 26))) + n = n // 26 - 1 + if n < 0: + break + return "".join(reversed(letters)) + + +async def values_get( + *, + access_token: str, + spreadsheet_id: str, + ranges: str, + fields: str | None = None, +) -> dict[str, Any]: + spreadsheet_id = extract_spreadsheet_id(spreadsheet_id) + params: dict[str, Any] = {} + # Sheets API rejects ranges="" as an invalid A1 range; omit when callers only want metadata. + if ranges: + params["ranges"] = ranges + if fields: + params["fields"] = fields + url = f"{SHEETS_API_BASE}/{spreadsheet_id}" + response = await _request_with_retry("GET", url, params=params, headers=_auth_headers(access_token)) + _raise_for_error(response) + return response.json() or {} + + +async def values_append( + *, + access_token: str, + spreadsheet_id: str, + range_: str, + values: list[list[Any]], + value_input_option: str = "USER_ENTERED", + insert_data_option: str = "INSERT_ROWS", +) -> dict[str, Any]: + spreadsheet_id = extract_spreadsheet_id(spreadsheet_id) + url = f"{SHEETS_API_BASE}/{spreadsheet_id}/values/{quote(range_, safe='')}:append" + params = {"valueInputOption": value_input_option, "insertDataOption": insert_data_option} + body = {"range": range_, "majorDimension": "ROWS", "values": values} + response = await _request_with_retry( + "POST", + url, + params=params, + json_body=body, + headers={**_auth_headers(access_token), "Content-Type": "application/json"}, + idempotent=False, + ) + _raise_for_error(response) + return response.json() or {} + + +async def values_update( + *, + access_token: str, + spreadsheet_id: str, + range_: str, + values: list[list[Any]], + value_input_option: str = "USER_ENTERED", +) -> dict[str, Any]: + spreadsheet_id = extract_spreadsheet_id(spreadsheet_id) + url = f"{SHEETS_API_BASE}/{spreadsheet_id}/values/{quote(range_, safe='')}" + params = {"valueInputOption": value_input_option} + body = {"range": range_, "majorDimension": "ROWS", "values": values} + response = await _request_with_retry( + "PUT", + url, + params=params, + json_body=body, + headers={**_auth_headers(access_token), "Content-Type": "application/json"}, + ) + _raise_for_error(response) + return response.json() or {} + + +async def batch_update( + *, + access_token: str, + spreadsheet_id: str, + requests: list[dict[str, Any]], +) -> dict[str, Any]: + spreadsheet_id = extract_spreadsheet_id(spreadsheet_id) + url = f"{SHEETS_API_BASE}/{spreadsheet_id}:batchUpdate" + response = await _request_with_retry( + "POST", + url, + json_body={"requests": requests}, + headers={**_auth_headers(access_token), "Content-Type": "application/json"}, + idempotent=False, + ) + _raise_for_error(response) + return response.json() or {} + + +async def get_sheet_id_by_title( + *, + access_token: str, + spreadsheet_id: str, + sheet_title: str, +) -> int | None: + payload = await values_get( + access_token=access_token, + spreadsheet_id=spreadsheet_id, + ranges="", + fields="sheets(properties(sheetId,title))", + ) + for sheet in payload.get("sheets") or []: + properties = sheet.get("properties") or {} + if str(properties.get("title")) == sheet_title: + value = properties.get("sheetId") + if isinstance(value, int): + return value + return None + + +def _grid_props_from_sheet(sheet: dict[str, Any]) -> SheetGridProperties | None: + props = sheet.get("properties") or {} + sheet_id = props.get("sheetId") + title = props.get("title") + grid = props.get("gridProperties") or {} + column_count = grid.get("columnCount") + row_count = grid.get("rowCount") + if ( + not isinstance(sheet_id, int) + or not isinstance(title, str) + or not isinstance(column_count, int) + or not isinstance(row_count, int) + ): + return None + return SheetGridProperties( + sheet_id=int(sheet_id), + title=title, + column_count=int(column_count), + row_count=int(row_count), + ) + + +async def get_sheet_grid_properties( + *, + access_token: str, + spreadsheet_id: str, + sheet_title: str, +) -> SheetGridProperties | None: + """Return the named tab's grid dimensions, or None if missing or malformed.""" + payload = await values_get( + access_token=access_token, + spreadsheet_id=spreadsheet_id, + ranges="", + fields="sheets(properties(sheetId,title,gridProperties(columnCount,rowCount)))", + ) + for sheet in payload.get("sheets") or []: + if str((sheet.get("properties") or {}).get("title")) != sheet_title: + continue + return _grid_props_from_sheet(sheet) + return None + + +async def get_sheet_grid_properties_by_id( + *, + access_token: str, + spreadsheet_id: str, + sheet_id: int, +) -> SheetGridProperties | None: + """Return the tab's grid dimensions matched by numeric sheetId, or None.""" + payload = await values_get( + access_token=access_token, + spreadsheet_id=spreadsheet_id, + ranges="", + fields="sheets(properties(sheetId,title,gridProperties(columnCount,rowCount)))", + ) + for sheet in payload.get("sheets") or []: + if (sheet.get("properties") or {}).get("sheetId") != sheet_id: + continue + return _grid_props_from_sheet(sheet) + return None + + +async def get_sheet_headers( + *, + access_token: str, + spreadsheet_id: str, + sheet_title: str, +) -> list[SheetHeader]: + """Return the non-blank cells of row 1 of ``sheet_title`` keyed by A1 column letter.""" + spreadsheet_id = extract_spreadsheet_id(spreadsheet_id) + a1 = f"{quote_sheet_name(sheet_title)}!A1:ZZZ1" + encoded = quote(a1, safe="") + url = f"{SHEETS_API_BASE}/{spreadsheet_id}/values/{encoded}" + response = await _request_with_retry("GET", url, headers=_auth_headers(access_token)) + _raise_for_error(response) + payload = response.json() or {} + values = payload.get("values") or [] + row = values[0] if values else [] + out: list[SheetHeader] = [] + for index, cell in enumerate(row): + name = str(cell).strip() if cell is not None else "" + if not name: + continue + out.append(SheetHeader(letter=_column_index_to_letter(index), name=name)) + return out diff --git a/tests/unit/force_stub_app.py b/tests/unit/force_stub_app.py index dc334f478..d115e1168 100644 --- a/tests/unit/force_stub_app.py +++ b/tests/unit/force_stub_app.py @@ -53,6 +53,7 @@ def create_forge_stub_app() -> ForgeApp: fake_app_module.AUTO_COMPLETION_LLM_API_HANDLER = AsyncMock() fake_app_module.EXPERIMENTATION_PROVIDER = _LazyNamespace() fake_app_module.STORAGE = _LazyNamespace() + fake_app_module.CACHE = _LazyNamespace() return fake_app_module diff --git a/tests/unit/google/__init__.py b/tests/unit/google/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/unit/google/test_agent_function_google_credentials.py b/tests/unit/google/test_agent_function_google_credentials.py new file mode 100644 index 000000000..7d3655a96 --- /dev/null +++ b/tests/unit/google/test_agent_function_google_credentials.py @@ -0,0 +1,159 @@ +"""Cover the OSS ``AgentFunction.get_google_*_credentials`` paths. + +These methods used to be no-ops; SKY-9463 wired them through the OSS +``google_oauth_service``. The tests below pin down the success and the +failure-modes (missing encryption, missing credential, refresh failure) +so a regression doesn't silently break Sheets blocks in OSS. +""" + +from types import SimpleNamespace +from unittest.mock import AsyncMock + +import pytest + +from skyvern.forge.agent_functions import AgentFunction +from skyvern.forge.sdk.services import google_oauth_service + + +@pytest.mark.asyncio +async def test_get_google_sheets_credentials_returns_token_on_success( + monkeypatch: pytest.MonkeyPatch, +) -> None: + secrets = google_oauth_service.GoogleCredentialSecrets(refresh_token="rt", scopes=[]) + monkeypatch.setattr( + google_oauth_service, + "load_credential_secrets", + AsyncMock(return_value=secrets), + ) + monkeypatch.setattr( + google_oauth_service, + "access_token_from_secrets", + AsyncMock(return_value="ya29.access-token"), + ) + + result = await AgentFunction().get_google_sheets_credentials( + organization_id="org_1", + credential_id="goac_1", + ) + assert result == "ya29.access-token" + + +@pytest.mark.asyncio +async def test_get_google_sheets_credentials_returns_none_when_encryption_disabled( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Missing ENABLE_ENCRYPTION must surface as None so callers route to reconnect.""" + + async def raise_encryption_not_configured(*_args, **_kwargs): + raise google_oauth_service.EncryptionNotConfiguredError("disabled") + + monkeypatch.setattr( + google_oauth_service, + "load_credential_secrets", + raise_encryption_not_configured, + ) + + result = await AgentFunction().get_google_sheets_credentials( + organization_id="org_1", + credential_id="goac_1", + ) + assert result is None + + +@pytest.mark.asyncio +async def test_get_google_sheets_credentials_returns_none_when_credential_missing( + monkeypatch: pytest.MonkeyPatch, +) -> None: + async def raise_missing(*_args, **_kwargs): + raise ValueError("No active Google OAuth credential found: goac_missing") + + monkeypatch.setattr(google_oauth_service, "load_credential_secrets", raise_missing) + + result = await AgentFunction().get_google_sheets_credentials( + organization_id="org_1", + credential_id="goac_missing", + ) + assert result is None + + +@pytest.mark.asyncio +async def test_get_google_sheets_credentials_returns_none_on_refresh_failure( + monkeypatch: pytest.MonkeyPatch, +) -> None: + secrets = google_oauth_service.GoogleCredentialSecrets(refresh_token="rt", scopes=[]) + monkeypatch.setattr( + google_oauth_service, + "load_credential_secrets", + AsyncMock(return_value=secrets), + ) + + async def raise_refresh_failure(*_args, **_kwargs): + raise google_oauth_service.MissingAccessTokenError("Google token refresh failed") + + monkeypatch.setattr(google_oauth_service, "access_token_from_secrets", raise_refresh_failure) + + result = await AgentFunction().get_google_sheets_credentials( + organization_id="org_1", + credential_id="goac_1", + ) + assert result is None + + +@pytest.mark.asyncio +async def test_get_google_workspace_credentials_returns_credentials_on_success( + monkeypatch: pytest.MonkeyPatch, +) -> None: + secrets = google_oauth_service.GoogleCredentialSecrets(refresh_token="rt", scopes=[]) + fake_credentials = SimpleNamespace(token="ya29.access-token") + monkeypatch.setattr( + google_oauth_service, + "load_credential_secrets", + AsyncMock(return_value=secrets), + ) + monkeypatch.setattr( + google_oauth_service, + "credentials_from_secrets", + AsyncMock(return_value=fake_credentials), + ) + + result = await AgentFunction().get_google_workspace_credentials( + organization_id="org_1", + credential_id="goac_1", + ) + assert result is fake_credentials + + +@pytest.mark.asyncio +async def test_get_google_workspace_credentials_returns_none_when_encryption_disabled( + monkeypatch: pytest.MonkeyPatch, +) -> None: + async def raise_encryption_not_configured(*_args, **_kwargs): + raise google_oauth_service.EncryptionNotConfiguredError("disabled") + + monkeypatch.setattr( + google_oauth_service, + "load_credential_secrets", + raise_encryption_not_configured, + ) + + result = await AgentFunction().get_google_workspace_credentials( + organization_id="org_1", + credential_id="goac_1", + ) + assert result is None + + +@pytest.mark.asyncio +async def test_get_google_workspace_credentials_returns_none_on_unexpected_error( + monkeypatch: pytest.MonkeyPatch, +) -> None: + async def raise_unexpected(*_args, **_kwargs): + raise RuntimeError("network blew up") + + monkeypatch.setattr(google_oauth_service, "load_credential_secrets", raise_unexpected) + + result = await AgentFunction().get_google_workspace_credentials( + organization_id="org_1", + credential_id="goac_1", + ) + assert result is None diff --git a/tests/unit/google/test_google_oauth.py b/tests/unit/google/test_google_oauth.py new file mode 100644 index 000000000..f8df0d2cd --- /dev/null +++ b/tests/unit/google/test_google_oauth.py @@ -0,0 +1,1003 @@ +from types import SimpleNamespace +from typing import Any +from unittest.mock import AsyncMock +from urllib.parse import parse_qs, urlparse + +import httpx +import pytest +from pydantic import ValidationError +from sqlalchemy.sql.elements import BindParameter + +from skyvern.forge.sdk.encrypt.base import EncryptMethod +from skyvern.forge.sdk.schemas.google_oauth import UpdateGoogleOAuthCredentialRequest +from skyvern.forge.sdk.services import google_oauth_service + + +def _unwrap_bind(value: Any) -> Any: + """SQLAlchemy wraps literal values passed to .values(...) in BindParameter; unwrap for equality checks.""" + return value.value if isinstance(value, BindParameter) else value + + +def _default_scopes_list() -> list[str]: + return list(google_oauth_service.GOOGLE_SHEETS_SCOPES) + + +def test_coerce_scopes_accepts_strings_and_iterables() -> None: + assert google_oauth_service._coerce_scopes("https://a/scope https://b/scope") == [ + "https://a/scope", + "https://b/scope", + ] + assert google_oauth_service._coerce_scopes("https://a/scope, https://b/scope") == [ + "https://a/scope", + "https://b/scope", + ] + assert google_oauth_service._coerce_scopes(["https://a", " https://b "]) == ["https://a", "https://b"] + assert google_oauth_service._coerce_scopes(None) == _default_scopes_list() + assert google_oauth_service._coerce_scopes("") == _default_scopes_list() + + +def test_google_sheets_scopes_includes_drive_file_and_metadata_readonly() -> None: + scopes = google_oauth_service.GOOGLE_SHEETS_SCOPES + assert "https://www.googleapis.com/auth/spreadsheets" in scopes + assert "https://www.googleapis.com/auth/drive.file" in scopes + assert "https://www.googleapis.com/auth/drive.metadata.readonly" in scopes + + +def test_sheets_api_runtime_defaults_match_previous_hardcoded_values() -> None: + """Sheets timeout/retry settings default to known values so unset envs + produce no behavior change for upgrading deployments.""" + from skyvern.config import Settings + + fresh = Settings() + assert fresh.GOOGLE_SHEETS_API_TIMEOUT_SECONDS == 30.0 + assert fresh.GOOGLE_SHEETS_API_MAX_RETRIES == 3 + + +def test_build_authorize_url_includes_required_params(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(google_oauth_service.settings, "GOOGLE_OAUTH_CLIENT_ID", "cid", raising=False) + monkeypatch.setattr(google_oauth_service.settings, "GOOGLE_OAUTH_CLIENT_SECRET", "csecret", raising=False) + monkeypatch.setattr(google_oauth_service.settings, "GOOGLE_OAUTH_REDIRECT_HOSTS", ["app"], raising=False) + + url, code_verifier = google_oauth_service.build_authorize_url( + redirect_uri="https://app/settings/google/callback", + state="abc123", + ) + + parsed = urlparse(url) + assert f"{parsed.scheme}://{parsed.netloc}{parsed.path}" == google_oauth_service.GOOGLE_AUTHORIZE_ENDPOINT + params = {k: v[0] for k, v in parse_qs(parsed.query).items()} + assert params["response_type"] == "code" + assert params["client_id"] == "cid" + assert params["redirect_uri"] == "https://app/settings/google/callback" + assert params["scope"] == " ".join(google_oauth_service.GOOGLE_SHEETS_SCOPES) + assert params["access_type"] == "offline" + assert params["prompt"] == "consent" + assert params["state"] == "abc123" + # PKCE: a code_challenge must be on the URL and the verifier returned for replay. + assert params["code_challenge_method"] == "S256" + assert params["code_challenge"] + assert code_verifier and len(code_verifier) >= 43 + + +def test_build_authorize_url_passes_autogenerate_code_verifier_explicitly(monkeypatch: pytest.MonkeyPatch) -> None: + """Explicit ``autogenerate_code_verifier=True`` so a library default flip can't silently drop PKCE.""" + monkeypatch.setattr(google_oauth_service.settings, "GOOGLE_OAUTH_CLIENT_ID", "cid", raising=False) + monkeypatch.setattr(google_oauth_service.settings, "GOOGLE_OAUTH_CLIENT_SECRET", "csecret", raising=False) + monkeypatch.setattr(google_oauth_service.settings, "GOOGLE_OAUTH_REDIRECT_HOSTS", ["x"], raising=False) + + captured: dict = {} + + class _FakeFlow: + code_verifier = "ver-fake" + + def authorization_url(self, **_kwargs): + return "https://accounts.google.com/o/oauth2/v2/auth", "state" + + @classmethod + def from_client_config(cls, *args, **kwargs): + captured["kwargs"] = kwargs + return cls() + + monkeypatch.setattr(google_oauth_service, "Flow", _FakeFlow) + + google_oauth_service.build_authorize_url(redirect_uri="https://x/cb", state="s") + + assert captured["kwargs"].get("autogenerate_code_verifier") is True + + +def test_build_authorize_url_without_client_id_raises(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(google_oauth_service.settings, "GOOGLE_OAUTH_CLIENT_ID", None, raising=False) + monkeypatch.setattr(google_oauth_service.settings, "GOOGLE_OAUTH_CLIENT_SECRET", "csecret", raising=False) + + with pytest.raises(ValueError, match="client credentials"): + google_oauth_service.build_authorize_url(redirect_uri="https://x", state="s") + + +def test_build_authorize_url_without_client_secret_raises(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(google_oauth_service.settings, "GOOGLE_OAUTH_CLIENT_ID", "cid", raising=False) + monkeypatch.setattr(google_oauth_service.settings, "GOOGLE_OAUTH_CLIENT_SECRET", None, raising=False) + + with pytest.raises(ValueError, match="client credentials"): + google_oauth_service.build_authorize_url(redirect_uri="https://x", state="s") + + +def test_build_authorize_url_rejects_redirect_uri_not_in_allowlist(monkeypatch: pytest.MonkeyPatch) -> None: + """Defense-in-depth: build_authorize_url must self-validate redirect_uri so direct callers + (outside start_authorization) cannot bypass the host allowlist.""" + monkeypatch.setattr(google_oauth_service.settings, "GOOGLE_OAUTH_CLIENT_ID", "cid", raising=False) + monkeypatch.setattr(google_oauth_service.settings, "GOOGLE_OAUTH_CLIENT_SECRET", "csecret", raising=False) + monkeypatch.setattr( + google_oauth_service.settings, + "GOOGLE_OAUTH_REDIRECT_HOSTS", + ["app.skyvern.com"], + raising=False, + ) + + with pytest.raises(google_oauth_service.InvalidRedirectURIError): + google_oauth_service.build_authorize_url( + redirect_uri="https://evil.example.com/callback", + state="abc", + ) + + +def test_build_authorize_url_rejects_http_for_non_loopback_host(monkeypatch: pytest.MonkeyPatch) -> None: + """Hostname-only allowlist plus an http URI must be rejected even when called directly.""" + monkeypatch.setattr(google_oauth_service.settings, "GOOGLE_OAUTH_CLIENT_ID", "cid", raising=False) + monkeypatch.setattr(google_oauth_service.settings, "GOOGLE_OAUTH_CLIENT_SECRET", "csecret", raising=False) + monkeypatch.setattr( + google_oauth_service.settings, + "GOOGLE_OAUTH_REDIRECT_HOSTS", + ["app.skyvern.com"], + raising=False, + ) + + with pytest.raises(google_oauth_service.InvalidRedirectURIError, match="https"): + google_oauth_service.build_authorize_url( + redirect_uri="http://app.skyvern.com/callback", + state="abc", + ) + + +@pytest.mark.asyncio +async def test_start_authorization_persists_verifier_and_returns_url(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(google_oauth_service.settings, "ENABLE_ENCRYPTION", True, raising=False) + monkeypatch.setattr(google_oauth_service.settings, "GOOGLE_OAUTH_CLIENT_ID", "cid", raising=False) + monkeypatch.setattr(google_oauth_service.settings, "GOOGLE_OAUTH_CLIENT_SECRET", "csecret", raising=False) + monkeypatch.setattr(google_oauth_service.settings, "GOOGLE_OAUTH_REDIRECT_HOSTS", ["x"], raising=False) + monkeypatch.setattr(google_oauth_service, "generate_google_oauth_credential_id", lambda: "goac_test") + + insert_mock = AsyncMock(return_value=SimpleNamespace(id="goac_test", organization_id="org_1")) + fake_repo = SimpleNamespace(insert_pending_credential=insert_mock) + monkeypatch.setattr(google_oauth_service.app, "DATABASE", SimpleNamespace(google_oauth=fake_repo), raising=False) + + result = await google_oauth_service.start_authorization( + organization_id="org_1", + redirect_uri="https://x/cb", + credential_name="my-cred", + ) + + assert result.authorize_url.startswith(google_oauth_service.GOOGLE_AUTHORIZE_ENDPOINT) + assert result.state + insert_mock.assert_awaited_once() + insert_kwargs = insert_mock.await_args.kwargs + assert insert_kwargs["organization_id"] == "org_1" + assert insert_kwargs["credential_name"] == "my-cred" + assert insert_kwargs["consent_redirect_uri"] == "https://x/cb" + assert insert_kwargs["consent_nonce"] == result.state + # The verifier must be in the same insert as everything else — no second + # round-trip — so a crash mid-flow can't leave a verifier-less pending row. + assert insert_kwargs["consent_code_verifier"] + assert len(insert_kwargs["consent_code_verifier"]) >= 43 + + +@pytest.mark.asyncio +async def test_start_authorization_refuses_without_encryption(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(google_oauth_service.settings, "ENABLE_ENCRYPTION", False, raising=False) + fake_repo = SimpleNamespace(insert_pending_credential=AsyncMock()) + monkeypatch.setattr(google_oauth_service.app, "DATABASE", SimpleNamespace(google_oauth=fake_repo), raising=False) + + with pytest.raises(google_oauth_service.EncryptionNotConfiguredError): + await google_oauth_service.start_authorization( + organization_id="org_1", + redirect_uri="https://x/cb", + ) + fake_repo.insert_pending_credential.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_promote_pending_credential_encrypts_and_calls_repo(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(google_oauth_service.settings, "ENABLE_ENCRYPTION", True, raising=False) + encrypt_mock = AsyncMock(return_value="ENC::rt") + monkeypatch.setattr(google_oauth_service, "encryptor", SimpleNamespace(encrypt=encrypt_mock)) + + promoted_schema = SimpleNamespace(id="goac_1", organization_id="org_1") + promote_mock = AsyncMock(return_value=promoted_schema) + fake_repo = SimpleNamespace(promote_pending_to_active=promote_mock) + monkeypatch.setattr(google_oauth_service.app, "DATABASE", SimpleNamespace(google_oauth=fake_repo), raising=False) + + result = await google_oauth_service.promote_pending_credential( + organization_id="org_1", + nonce="nonce-xyz", + refresh_token="rt-plain", + scopes_granted="https://a https://b", + ) + + assert result is promoted_schema + encrypt_mock.assert_awaited_once_with("rt-plain", EncryptMethod.AES) + promote_mock.assert_awaited_once() + kwargs = promote_mock.await_args.kwargs + assert kwargs["organization_id"] == "org_1" + assert kwargs["nonce"] == "nonce-xyz" + assert kwargs["encrypted_refresh_token"] == "ENC::rt" + assert kwargs["encrypted_method"] == EncryptMethod.AES + assert kwargs["scopes_granted"] == ["https://a", "https://b"] + + +@pytest.mark.asyncio +async def test_load_pending_consent_context_delegates_to_repo(monkeypatch: pytest.MonkeyPatch) -> None: + from skyvern.forge.sdk.db.repositories.google_oauth import PendingConsentContext + + expected = PendingConsentContext( + credential_id="goac_1", + consent_redirect_uri="https://x/cb", + consent_code_verifier="ver-abc", + ) + fake_repo = SimpleNamespace(load_pending_by_nonce=AsyncMock(return_value=expected)) + monkeypatch.setattr(google_oauth_service.app, "DATABASE", SimpleNamespace(google_oauth=fake_repo), raising=False) + + result = await google_oauth_service.load_pending_consent_context( + organization_id="org_1", + nonce="nonce-xyz", + ) + assert result is expected + fake_repo.load_pending_by_nonce.assert_awaited_once_with(organization_id="org_1", nonce="nonce-xyz") + + +class _FakeOAuth2Session: + def __init__(self, token: dict | None = None) -> None: + self.token = token or {} + + +class _FakeFlow: + """Mirrors the real google-auth-oauthlib Flow surface exchange_code_for_tokens touches.""" + + def __init__(self, credentials: Any, session_token: dict | None, captured: dict) -> None: + self.credentials = credentials + self.oauth2session = _FakeOAuth2Session(token=session_token) + self._captured = captured + + def fetch_token(self, code: str, code_verifier: str | None = None) -> None: + self._captured["code"] = code + self._captured["fetch_token_code_verifier"] = code_verifier + + +def _install_fake_flow( + monkeypatch: pytest.MonkeyPatch, + *, + credentials: Any, + session_token: dict | None, +) -> dict: + captured: dict = {} + + class _FlowFactory: + @classmethod + def from_client_config( + cls, client_config: dict, scopes=None, redirect_uri=None, state=None, code_verifier=None + ) -> _FakeFlow: + captured["client_config"] = client_config + captured["scopes"] = scopes + captured["redirect_uri"] = redirect_uri + captured["init_code_verifier"] = code_verifier + return _FakeFlow(credentials=credentials, session_token=session_token, captured=captured) + + monkeypatch.setattr(google_oauth_service, "Flow", _FlowFactory) + return captured + + +@pytest.mark.asyncio +async def test_exchange_code_for_tokens_uses_granted_scopes(monkeypatch: pytest.MonkeyPatch) -> None: + """Real google-auth leaves Credentials.scopes=None when Flow is built with scopes=None; + the granted scope lives on Credentials.granted_scopes (passed through from the token response).""" + monkeypatch.setattr(google_oauth_service.settings, "GOOGLE_OAUTH_CLIENT_ID", "cid", raising=False) + monkeypatch.setattr(google_oauth_service.settings, "GOOGLE_OAUTH_CLIENT_SECRET", "secret", raising=False) + + creds = SimpleNamespace( + token="at-from-flow", + refresh_token="rt-from-flow", + scopes=None, + granted_scopes="https://a/scope https://b/scope", + expiry=None, + ) + captured = _install_fake_flow( + monkeypatch, + credentials=creds, + session_token={"scope": "https://a/scope https://b/scope", "access_token": "at-from-flow"}, + ) + + result = await google_oauth_service.exchange_code_for_tokens( + code="abc", redirect_uri="https://x/cb", code_verifier="ver-xyz" + ) + + assert result == { + "access_token": "at-from-flow", + "refresh_token": "rt-from-flow", + "scope": "https://a/scope https://b/scope", + "expiry": None, + } + assert captured["code"] == "abc" + assert captured["redirect_uri"] == "https://x/cb" + assert captured["client_config"]["web"]["client_id"] == "cid" + assert captured["client_config"]["web"]["client_secret"] == "secret" + # ``Flow.from_client_config`` ignores ``code_verifier`` — PKCE actually replays + # via ``fetch_token``, so only the fetch-time verifier matters. + assert captured["fetch_token_code_verifier"] == "ver-xyz" + + +@pytest.mark.asyncio +async def test_exchange_code_for_tokens_accepts_granted_scopes_sequence(monkeypatch: pytest.MonkeyPatch) -> None: + """Future-proof: if upstream starts handing back granted_scopes as a list, we must still serialize it.""" + monkeypatch.setattr(google_oauth_service.settings, "GOOGLE_OAUTH_CLIENT_ID", "cid", raising=False) + monkeypatch.setattr(google_oauth_service.settings, "GOOGLE_OAUTH_CLIENT_SECRET", "secret", raising=False) + + creds = SimpleNamespace( + token="at", + refresh_token="rt", + scopes=None, + granted_scopes=["https://a/scope", "https://b/scope"], + expiry=None, + ) + _install_fake_flow(monkeypatch, credentials=creds, session_token={"scope": "", "access_token": "at"}) + + result = await google_oauth_service.exchange_code_for_tokens( + code="abc", redirect_uri="https://x/cb", code_verifier="v" + ) + assert result["scope"] == "https://a/scope https://b/scope" + + +@pytest.mark.asyncio +async def test_exchange_code_for_tokens_falls_back_to_session_token_scope(monkeypatch: pytest.MonkeyPatch) -> None: + """Belt-and-suspenders: if a library variant leaves granted_scopes empty, read the raw token response.""" + monkeypatch.setattr(google_oauth_service.settings, "GOOGLE_OAUTH_CLIENT_ID", "cid", raising=False) + monkeypatch.setattr(google_oauth_service.settings, "GOOGLE_OAUTH_CLIENT_SECRET", "secret", raising=False) + + creds = SimpleNamespace( + token="at", + refresh_token="rt", + scopes=None, + granted_scopes=None, + expiry=None, + ) + _install_fake_flow( + monkeypatch, + credentials=creds, + session_token={"scope": "https://a/scope", "access_token": "at"}, + ) + + result = await google_oauth_service.exchange_code_for_tokens( + code="abc", redirect_uri="https://x/cb", code_verifier="v" + ) + assert result["scope"] == "https://a/scope" + + +@pytest.mark.asyncio +async def test_exchange_code_raises_without_client_creds(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(google_oauth_service.settings, "GOOGLE_OAUTH_CLIENT_ID", None, raising=False) + monkeypatch.setattr(google_oauth_service.settings, "GOOGLE_OAUTH_CLIENT_SECRET", None, raising=False) + + with pytest.raises(ValueError, match="client credentials"): + await google_oauth_service.exchange_code_for_tokens("code", "https://x/cb", code_verifier="v") + + +@pytest.mark.asyncio +async def test_revoke_credential_loads_decrypts_revokes_upstream_marks(monkeypatch: pytest.MonkeyPatch) -> None: + from skyvern.forge.sdk.db.repositories.google_oauth import RevocableCiphertext + + load_mock = AsyncMock( + return_value=RevocableCiphertext( + exists=True, + encrypted_refresh_token="ENC::token", + encrypted_method=EncryptMethod.AES, + ) + ) + mark_mock = AsyncMock(return_value="goac_1") + monkeypatch.setattr( + google_oauth_service.app, + "DATABASE", + SimpleNamespace( + google_oauth=SimpleNamespace( + load_ciphertext_for_revoke=load_mock, + mark_revoked_and_scrub=mark_mock, + ) + ), + raising=False, + ) + + decrypt_mock = AsyncMock(return_value="refresh-123") + monkeypatch.setattr(google_oauth_service, "encryptor", SimpleNamespace(decrypt=decrypt_mock)) + upstream_mock = AsyncMock() + monkeypatch.setattr(google_oauth_service, "_revoke_refresh_token_at_google", upstream_mock) + + fake_cache = SimpleNamespace(set=AsyncMock()) + monkeypatch.setattr(google_oauth_service.app, "CACHE", fake_cache, raising=False) + + revoked = await google_oauth_service.revoke_credential(organization_id="org_1", credential_id="goac_1") + assert revoked is True + decrypt_mock.assert_awaited_once_with("ENC::token", EncryptMethod.AES) + upstream_mock.assert_awaited_once_with("refresh-123") + mark_mock.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_revoke_credential_missing_returns_false(monkeypatch: pytest.MonkeyPatch) -> None: + from skyvern.forge.sdk.db.repositories.google_oauth import RevocableCiphertext + + load_mock = AsyncMock(return_value=RevocableCiphertext(exists=False)) + upstream_mock = AsyncMock() + mark_mock = AsyncMock() + monkeypatch.setattr( + google_oauth_service.app, + "DATABASE", + SimpleNamespace( + google_oauth=SimpleNamespace( + load_ciphertext_for_revoke=load_mock, + mark_revoked_and_scrub=mark_mock, + ) + ), + raising=False, + ) + monkeypatch.setattr(google_oauth_service, "_revoke_refresh_token_at_google", upstream_mock) + + revoked = await google_oauth_service.revoke_credential(organization_id="o", credential_id="c") + assert revoked is False + upstream_mock.assert_not_awaited() + mark_mock.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_revoke_google_endpoint_swallows_errors(monkeypatch: pytest.MonkeyPatch) -> None: + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response(500, text="boom") + + transport = httpx.MockTransport(handler) + real_client = httpx.AsyncClient + monkeypatch.setattr( + google_oauth_service.httpx, + "AsyncClient", + lambda *a, **kw: real_client(transport=transport), + ) + + await google_oauth_service._revoke_refresh_token_at_google("rt") + + +@pytest.mark.asyncio +async def test_load_credential_secrets_decrypts_repo_payload(monkeypatch: pytest.MonkeyPatch) -> None: + from skyvern.forge.sdk.db.repositories.google_oauth import ActiveCredentialCiphertext + + payload = ActiveCredentialCiphertext( + encrypted_refresh_token="ENC::rt", + encrypted_method=EncryptMethod.AES, + scopes_granted=["https://a", "https://b"], + ) + fake_repo = SimpleNamespace(load_active_ciphertext=AsyncMock(return_value=payload)) + monkeypatch.setattr(google_oauth_service.app, "DATABASE", SimpleNamespace(google_oauth=fake_repo), raising=False) + + decrypt_mock = AsyncMock(return_value="refresh-123") + monkeypatch.setattr(google_oauth_service, "encryptor", SimpleNamespace(decrypt=decrypt_mock)) + + secrets = await google_oauth_service.load_credential_secrets( + organization_id="org_1", + credential_id="goac_1", + ) + + assert secrets.refresh_token == "refresh-123" + assert secrets.scopes == ["https://a", "https://b"] + decrypt_mock.assert_awaited_once_with("ENC::rt", EncryptMethod.AES) + + +@pytest.mark.asyncio +async def test_load_credential_secrets_raises_when_missing(monkeypatch: pytest.MonkeyPatch) -> None: + fake_repo = SimpleNamespace(load_active_ciphertext=AsyncMock(return_value=None)) + monkeypatch.setattr(google_oauth_service.app, "DATABASE", SimpleNamespace(google_oauth=fake_repo), raising=False) + + with pytest.raises(ValueError, match="No active Google OAuth credential"): + await google_oauth_service.load_credential_secrets( + organization_id="org_1", + credential_id="goac_missing", + ) + + +@pytest.mark.asyncio +async def test_get_credentials_for_org_delegates_to_repo(monkeypatch: pytest.MonkeyPatch) -> None: + creds = [SimpleNamespace(id="goac_1"), SimpleNamespace(id="goac_2")] + list_mock = AsyncMock(return_value=creds) + monkeypatch.setattr( + google_oauth_service.app, + "DATABASE", + SimpleNamespace(google_oauth=SimpleNamespace(list_active_for_org=list_mock)), + raising=False, + ) + + result = await google_oauth_service.get_credentials_for_org(organization_id="org_1") + assert result is creds + list_mock.assert_awaited_once_with(organization_id="org_1") + + +@pytest.mark.asyncio +async def test_access_token_from_secrets_calls_refresh(monkeypatch: pytest.MonkeyPatch) -> None: + refresh_mock = AsyncMock(return_value={"access_token": "at-456"}) + monkeypatch.setattr(google_oauth_service, "refresh_access_token", refresh_mock) + + secrets = google_oauth_service.GoogleCredentialSecrets( + refresh_token="rt-1", + scopes=["https://a"], + ) + + token = await google_oauth_service.access_token_from_secrets(secrets) + + assert token == "at-456" + refresh_mock.assert_awaited_once_with("rt-1") + + +@pytest.mark.asyncio +async def test_access_token_from_secrets_missing_access_token_raises(monkeypatch: pytest.MonkeyPatch) -> None: + refresh_mock = AsyncMock(return_value={"scope": "foo"}) + monkeypatch.setattr(google_oauth_service, "refresh_access_token", refresh_mock) + + secrets = google_oauth_service.GoogleCredentialSecrets(refresh_token="rt", scopes=[]) + + with pytest.raises(google_oauth_service.MissingAccessTokenError): + await google_oauth_service.access_token_from_secrets(secrets) + + +def test_validate_redirect_uri_allowlist(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr( + google_oauth_service.settings, + "GOOGLE_OAUTH_REDIRECT_HOSTS", + ["app.skyvern.com"], + raising=False, + ) + + google_oauth_service._validate_redirect_uri("https://app.skyvern.com/google/callback") + + with pytest.raises(google_oauth_service.InvalidRedirectURIError): + google_oauth_service._validate_redirect_uri("https://evil.example.com/callback") + + +def test_validate_redirect_uri_empty_allowlist_dev_fallback(monkeypatch: pytest.MonkeyPatch) -> None: + """Empty allowlist + no client_id = dev fallback; keeps local dev ergonomic.""" + monkeypatch.setattr(google_oauth_service.settings, "GOOGLE_OAUTH_REDIRECT_HOSTS", [], raising=False) + monkeypatch.setattr(google_oauth_service.settings, "GOOGLE_OAUTH_CLIENT_ID", None, raising=False) + google_oauth_service._validate_redirect_uri("https://anywhere.example.com/cb") + + +def test_validate_redirect_uri_empty_allowlist_raises_when_client_id_set( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Empty allowlist + client_id configured = misconfigured prod; must fail closed.""" + monkeypatch.setattr(google_oauth_service.settings, "GOOGLE_OAUTH_REDIRECT_HOSTS", [], raising=False) + monkeypatch.setattr(google_oauth_service.settings, "GOOGLE_OAUTH_CLIENT_ID", "cid", raising=False) + with pytest.raises(google_oauth_service.InvalidRedirectURIError): + google_oauth_service._validate_redirect_uri("https://app.skyvern.com/cb") + + +def test_validate_redirect_uri_rejects_http_for_non_loopback(monkeypatch: pytest.MonkeyPatch) -> None: + """Hostname-only allowlist plus an http URI is the bypass claude[bot] flagged: + an attacker on http://allowed-host.com:9999 should not satisfy a check intended + for https://allowed-host.com.""" + monkeypatch.setattr( + google_oauth_service.settings, + "GOOGLE_OAUTH_REDIRECT_HOSTS", + ["app.skyvern.com"], + raising=False, + ) + with pytest.raises(google_oauth_service.InvalidRedirectURIError, match="https"): + google_oauth_service._validate_redirect_uri("http://app.skyvern.com/cb") + + +def test_validate_redirect_uri_allows_http_for_loopback(monkeypatch: pytest.MonkeyPatch) -> None: + """Local dev needs http://localhost; loopback hosts are exempted from the https requirement.""" + monkeypatch.setattr( + google_oauth_service.settings, + "GOOGLE_OAUTH_REDIRECT_HOSTS", + ["localhost", "127.0.0.1"], + raising=False, + ) + google_oauth_service._validate_redirect_uri("http://localhost:5173/cb") + google_oauth_service._validate_redirect_uri("http://127.0.0.1:8080/cb") + + +def test_validate_redirect_uri_normalizes_allowlist_case(monkeypatch: pytest.MonkeyPatch) -> None: + """``urlparse().hostname`` lowercases the URI host; the allowlist is lowercased + too so an operator who configures mixed-case entries doesn't reject every redirect.""" + monkeypatch.setattr( + google_oauth_service.settings, + "GOOGLE_OAUTH_REDIRECT_HOSTS", + ["MyApp.Example.Com"], + raising=False, + ) + google_oauth_service._validate_redirect_uri("https://myapp.example.com/cb") + google_oauth_service._validate_redirect_uri("https://MYAPP.EXAMPLE.COM/cb") + + +@pytest.mark.asyncio +async def test_refresh_access_token_uses_credentials_refresh(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(google_oauth_service.settings, "GOOGLE_OAUTH_CLIENT_ID", "cid", raising=False) + monkeypatch.setattr(google_oauth_service.settings, "GOOGLE_OAUTH_CLIENT_SECRET", "secret", raising=False) + + captured: dict = {} + + class _FakeCreds: + def __init__(self, **kwargs) -> None: + captured["init_kwargs"] = kwargs + self.token: str | None = None + self.expiry = None + + def refresh(self, request) -> None: + captured["refresh_request_type"] = type(request).__name__ + self.token = "at-refreshed" + + monkeypatch.setattr(google_oauth_service, "Credentials", _FakeCreds) + + result = await google_oauth_service.refresh_access_token("rt-1") + + assert result == {"access_token": "at-refreshed", "expiry": None} + assert captured["init_kwargs"]["refresh_token"] == "rt-1" + assert captured["init_kwargs"]["client_id"] == "cid" + assert captured["init_kwargs"]["token_uri"] == google_oauth_service.GOOGLE_TOKEN_ENDPOINT + assert captured["refresh_request_type"] == "Request" + + +@pytest.mark.asyncio +async def test_refresh_access_token_wraps_google_auth_error(monkeypatch: pytest.MonkeyPatch) -> None: + from google.auth.exceptions import RefreshError + + monkeypatch.setattr(google_oauth_service.settings, "GOOGLE_OAUTH_CLIENT_ID", "cid", raising=False) + monkeypatch.setattr(google_oauth_service.settings, "GOOGLE_OAUTH_CLIENT_SECRET", "secret", raising=False) + + class _FakeCreds: + def __init__(self, **_kwargs) -> None: + self.token = None + self.expiry = None + + def refresh(self, _request) -> None: + raise RefreshError("invalid_grant") + + monkeypatch.setattr(google_oauth_service, "Credentials", _FakeCreds) + + with pytest.raises(google_oauth_service.MissingAccessTokenError, match="refresh failed"): + await google_oauth_service.refresh_access_token("rt-bad") + + +@pytest.mark.asyncio +async def test_credentials_from_secrets_wraps_google_auth_error(monkeypatch: pytest.MonkeyPatch) -> None: + from google.auth.exceptions import RefreshError + + monkeypatch.setattr(google_oauth_service.settings, "GOOGLE_OAUTH_CLIENT_ID", "cid", raising=False) + monkeypatch.setattr(google_oauth_service.settings, "GOOGLE_OAUTH_CLIENT_SECRET", "secret", raising=False) + + class _FakeCreds: + def __init__(self, **_kwargs) -> None: + self.token = None + + def refresh(self, _request) -> None: + raise RefreshError("token revoked") + + monkeypatch.setattr(google_oauth_service, "Credentials", _FakeCreds) + + secrets = google_oauth_service.GoogleCredentialSecrets(refresh_token="rt-1", scopes=["https://a"]) + with pytest.raises(google_oauth_service.MissingAccessTokenError, match="refresh failed"): + await google_oauth_service.credentials_from_secrets(secrets) + + +@pytest.mark.asyncio +async def test_credentials_from_secrets_returns_refreshed_credentials(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(google_oauth_service.settings, "GOOGLE_OAUTH_CLIENT_ID", "cid", raising=False) + monkeypatch.setattr(google_oauth_service.settings, "GOOGLE_OAUTH_CLIENT_SECRET", "secret", raising=False) + + captured: dict = {} + + class _FakeCreds: + def __init__(self, **kwargs) -> None: + captured["init_kwargs"] = kwargs + self.token: str | None = None + + def refresh(self, request) -> None: + self.token = "at-refreshed" + + monkeypatch.setattr(google_oauth_service, "Credentials", _FakeCreds) + + secrets = google_oauth_service.GoogleCredentialSecrets( + refresh_token="rt-decoded", + scopes=["https://a", "https://b"], + ) + creds = await google_oauth_service.credentials_from_secrets(secrets) + + assert creds.token == "at-refreshed" + assert captured["init_kwargs"]["refresh_token"] == "rt-decoded" + assert captured["init_kwargs"]["scopes"] == ["https://a", "https://b"] + + +def test_update_google_oauth_credential_request_strips_whitespace() -> None: + request = UpdateGoogleOAuthCredentialRequest(credential_name=" Personal Gmail ") + assert request.credential_name == "Personal Gmail" + + +def test_update_google_oauth_credential_request_rejects_blank() -> None: + with pytest.raises(ValidationError): + UpdateGoogleOAuthCredentialRequest(credential_name=" ") + + +def test_update_google_oauth_credential_request_rejects_empty() -> None: + with pytest.raises(ValidationError): + UpdateGoogleOAuthCredentialRequest(credential_name="") + + +def test_update_google_oauth_credential_request_enforces_max_length() -> None: + with pytest.raises(ValidationError): + UpdateGoogleOAuthCredentialRequest(credential_name="x" * 129) + + +@pytest.mark.asyncio +async def test_rename_credential_delegates_to_repo(monkeypatch: pytest.MonkeyPatch) -> None: + renamed = SimpleNamespace(id="goac_1", credential_name="Personal Gmail") + rename_mock = AsyncMock(return_value=renamed) + monkeypatch.setattr( + google_oauth_service.app, + "DATABASE", + SimpleNamespace(google_oauth=SimpleNamespace(rename_active=rename_mock)), + raising=False, + ) + + result = await google_oauth_service.rename_credential( + organization_id="org_1", + credential_id="goac_1", + credential_name="Personal Gmail", + ) + assert result is renamed + rename_mock.assert_awaited_once() + kwargs = rename_mock.await_args.kwargs + assert kwargs["organization_id"] == "org_1" + assert kwargs["credential_id"] == "goac_1" + assert kwargs["credential_name"] == "Personal Gmail" + + +@pytest.mark.asyncio +async def test_rename_credential_returns_none_when_not_found(monkeypatch: pytest.MonkeyPatch) -> None: + rename_mock = AsyncMock(return_value=None) + monkeypatch.setattr( + google_oauth_service.app, + "DATABASE", + SimpleNamespace(google_oauth=SimpleNamespace(rename_active=rename_mock)), + raising=False, + ) + result = await google_oauth_service.rename_credential( + organization_id="org_1", + credential_id="goac_missing", + credential_name="Anything", + ) + assert result is None + + +def test_require_scopes_from_token_returns_scope_when_present() -> None: + from skyvern.forge.sdk.routes import google_oauth as oauth_route + + assert oauth_route._require_scopes_from_token({"scope": "https://a https://b"}) == ["https://a", "https://b"] + + +def test_require_scopes_from_token_raises_on_missing_scope() -> None: + from fastapi import HTTPException + + from skyvern.forge.sdk.routes import google_oauth as oauth_route + + with pytest.raises(HTTPException) as excinfo: + oauth_route._require_scopes_from_token({}) + assert excinfo.value.status_code == 400 + assert "scope" in excinfo.value.detail.lower() + + +def test_require_scopes_from_token_raises_on_empty_scope_string() -> None: + """Google returns an empty scope string on partial consent; must fail closed, not default.""" + from fastapi import HTTPException + + from skyvern.forge.sdk.routes import google_oauth as oauth_route + + with pytest.raises(HTTPException) as excinfo: + oauth_route._require_scopes_from_token({"scope": ""}) + assert excinfo.value.status_code == 400 + assert "scope" in excinfo.value.detail.lower() + + +def test_require_scopes_from_token_raises_on_whitespace_only_scope() -> None: + from fastapi import HTTPException + + from skyvern.forge.sdk.routes import google_oauth as oauth_route + + with pytest.raises(HTTPException): + oauth_route._require_scopes_from_token({"scope": " "}) + + +# --------------------------------------------------------------------------- +# _validate_app_origin tests +# --------------------------------------------------------------------------- + + +def test_validate_app_origin_exact_match_https(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr( + google_oauth_service.settings, + "GOOGLE_OAUTH_APP_ORIGINS", + ["https://app.skyvern.com"], + raising=False, + ) + google_oauth_service._validate_app_origin("https://app.skyvern.com") + + +def test_validate_app_origin_exact_match_http_localhost(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr( + google_oauth_service.settings, + "GOOGLE_OAUTH_APP_ORIGINS", + ["http://localhost:5173"], + raising=False, + ) + google_oauth_service._validate_app_origin("http://localhost:5173") + + +def test_validate_app_origin_rejects_non_matching_exact(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr( + google_oauth_service.settings, + "GOOGLE_OAUTH_APP_ORIGINS", + ["https://app.skyvern.com"], + raising=False, + ) + with pytest.raises(google_oauth_service.InvalidAppOriginError, match="not allowed"): + google_oauth_service._validate_app_origin("https://evil.example.com") + + +def test_validate_app_origin_suffix_wildcard_matches(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr( + google_oauth_service.settings, + "GOOGLE_OAUTH_APP_ORIGINS", + ["*.vercel.app"], + raising=False, + ) + google_oauth_service._validate_app_origin("https://skyvern-cloud-git-main-skyvern.vercel.app") + + +def test_validate_app_origin_suffix_wildcard_matches_with_port(monkeypatch: pytest.MonkeyPatch) -> None: + """Suffix matching strips the port — ``*.vercel.app`` accepts ``myapp.vercel.app:3000``.""" + monkeypatch.setattr( + google_oauth_service.settings, + "GOOGLE_OAUTH_APP_ORIGINS", + ["*.vercel.app"], + raising=False, + ) + google_oauth_service._validate_app_origin("https://myapp.vercel.app:3000") + + +def test_validate_app_origin_suffix_wildcard_rejects_bare_suffix(monkeypatch: pytest.MonkeyPatch) -> None: + """'vercel.app' itself must not match '*.vercel.app'.""" + monkeypatch.setattr( + google_oauth_service.settings, + "GOOGLE_OAUTH_APP_ORIGINS", + ["*.vercel.app"], + raising=False, + ) + with pytest.raises(google_oauth_service.InvalidAppOriginError): + google_oauth_service._validate_app_origin("https://vercel.app") + + +def test_validate_app_origin_suffix_wildcard_rejects_spoof_prefix(monkeypatch: pytest.MonkeyPatch) -> None: + """'attacker-vercel.app' must not match '*.vercel.app'.""" + monkeypatch.setattr( + google_oauth_service.settings, + "GOOGLE_OAUTH_APP_ORIGINS", + ["*.vercel.app"], + raising=False, + ) + with pytest.raises(google_oauth_service.InvalidAppOriginError): + google_oauth_service._validate_app_origin("https://attacker-vercel.app") + + +def test_validate_app_origin_suffix_wildcard_rejects_spoof_suffix(monkeypatch: pytest.MonkeyPatch) -> None: + """'vercel.app.evil.com' must not match '*.vercel.app'.""" + monkeypatch.setattr( + google_oauth_service.settings, + "GOOGLE_OAUTH_APP_ORIGINS", + ["*.vercel.app"], + raising=False, + ) + with pytest.raises(google_oauth_service.InvalidAppOriginError): + google_oauth_service._validate_app_origin("https://vercel.app.evil.com") + + +def test_validate_app_origin_wildcard_requires_https(monkeypatch: pytest.MonkeyPatch) -> None: + """Wildcard entries only match https, not http.""" + monkeypatch.setattr( + google_oauth_service.settings, + "GOOGLE_OAUTH_APP_ORIGINS", + ["*.vercel.app"], + raising=False, + ) + with pytest.raises(google_oauth_service.InvalidAppOriginError): + google_oauth_service._validate_app_origin("http://skyvern-cloud-git-main-skyvern.vercel.app") + + +def test_validate_app_origin_empty_allowlist_raises(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(google_oauth_service.settings, "GOOGLE_OAUTH_APP_ORIGINS", [], raising=False) + with pytest.raises(google_oauth_service.InvalidAppOriginError, match="not configured"): + google_oauth_service._validate_app_origin("https://app.skyvern.com") + + +@pytest.mark.asyncio +async def test_start_authorization_persists_app_origin(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(google_oauth_service.settings, "ENABLE_ENCRYPTION", True, raising=False) + monkeypatch.setattr(google_oauth_service.settings, "GOOGLE_OAUTH_CLIENT_ID", "cid", raising=False) + monkeypatch.setattr(google_oauth_service.settings, "GOOGLE_OAUTH_CLIENT_SECRET", "csecret", raising=False) + monkeypatch.setattr( + google_oauth_service.settings, "GOOGLE_OAUTH_REDIRECT_HOSTS", ["app-staging.skyvern.com"], raising=False + ) + monkeypatch.setattr(google_oauth_service.settings, "GOOGLE_OAUTH_APP_ORIGINS", ["*.vercel.app"], raising=False) + monkeypatch.setattr(google_oauth_service, "generate_google_oauth_credential_id", lambda: "goac_test2") + + insert_mock = AsyncMock(return_value=SimpleNamespace(id="goac_test2", organization_id="org_1")) + fake_repo = SimpleNamespace(insert_pending_credential=insert_mock) + monkeypatch.setattr(google_oauth_service.app, "DATABASE", SimpleNamespace(google_oauth=fake_repo), raising=False) + + result = await google_oauth_service.start_authorization( + organization_id="org_1", + redirect_uri="https://app-staging.skyvern.com/integrations/google/callback", + credential_name="Test", + app_origin="https://skyvern-cloud-git-branch-skyvern.vercel.app", + ) + + assert result.state + insert_mock.assert_awaited_once() + insert_kwargs = insert_mock.await_args.kwargs + assert insert_kwargs["consent_app_origin"] == "https://skyvern-cloud-git-branch-skyvern.vercel.app" + + +@pytest.mark.asyncio +async def test_start_authorization_rejects_bad_app_origin(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(google_oauth_service.settings, "ENABLE_ENCRYPTION", True, raising=False) + monkeypatch.setattr(google_oauth_service.settings, "GOOGLE_OAUTH_CLIENT_ID", "cid", raising=False) + monkeypatch.setattr(google_oauth_service.settings, "GOOGLE_OAUTH_CLIENT_SECRET", "csecret", raising=False) + monkeypatch.setattr( + google_oauth_service.settings, "GOOGLE_OAUTH_REDIRECT_HOSTS", ["app-staging.skyvern.com"], raising=False + ) + monkeypatch.setattr( + google_oauth_service.settings, "GOOGLE_OAUTH_APP_ORIGINS", ["https://app.skyvern.com"], raising=False + ) + + insert_mock = AsyncMock() + fake_repo = SimpleNamespace(insert_pending_credential=insert_mock) + monkeypatch.setattr(google_oauth_service.app, "DATABASE", SimpleNamespace(google_oauth=fake_repo), raising=False) + + with pytest.raises(google_oauth_service.InvalidAppOriginError): + await google_oauth_service.start_authorization( + organization_id="org_1", + redirect_uri="https://app-staging.skyvern.com/integrations/google/callback", + app_origin="https://evil.example.com", + ) + insert_mock.assert_not_awaited() + + +def test_google_oauth_credential_response_exposes_app_origin() -> None: + import datetime + + from skyvern.forge.sdk.schemas.google_oauth import GoogleOAuthCredentialBase, GoogleOAuthCredentialResponse + + cred = GoogleOAuthCredentialBase( + id="goac_1", + organization_id="o_1", + credential_name="Default", + provider="google", + state="active", + scopes_requested=[], + scopes_granted=[], + created_at=datetime.datetime.utcnow(), + modified_at=datetime.datetime.utcnow(), + ) + resp = GoogleOAuthCredentialResponse(credential=cred, app_origin="https://foo.vercel.app") + assert resp.app_origin == "https://foo.vercel.app" + + resp_no_origin = GoogleOAuthCredentialResponse(credential=cred) + assert resp_no_origin.app_origin is None diff --git a/tests/unit/google/test_google_oauth_repository.py b/tests/unit/google/test_google_oauth_repository.py new file mode 100644 index 000000000..fc492ee02 --- /dev/null +++ b/tests/unit/google/test_google_oauth_repository.py @@ -0,0 +1,197 @@ +import datetime +from typing import AsyncGenerator + +import pytest +import pytest_asyncio +from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine + +from skyvern.forge.sdk.db.base_alchemy_db import BaseAlchemyDB +from skyvern.forge.sdk.db.models import Base, GoogleOAuthCredentialModel # noqa: F401 - registers model on Base +from skyvern.forge.sdk.db.repositories.google_oauth import ( + STATE_PENDING_CONSENT, + GoogleOAuthRepository, +) +from skyvern.forge.sdk.encrypt.base import EncryptMethod +from skyvern.forge.sdk.schemas.google_oauth import GoogleOAuthCredentialBase + + +@pytest_asyncio.fixture +async def engine() -> AsyncGenerator[AsyncEngine, None]: + eng = create_async_engine("sqlite+aiosqlite:///:memory:") + async with eng.begin() as conn: + await conn.run_sync(Base.metadata.create_all) + yield eng + await eng.dispose() + + +@pytest_asyncio.fixture +async def repo(engine: AsyncEngine) -> GoogleOAuthRepository: + db = BaseAlchemyDB(engine) + return GoogleOAuthRepository(db.Session, debug_enabled=False) + + +@pytest.mark.asyncio +async def test_insert_pending_credential_returns_schema_without_greenlet_error( + repo: GoogleOAuthRepository, +) -> None: + expires_at = datetime.datetime.utcnow() + datetime.timedelta(minutes=10) + result = await repo.insert_pending_credential( + credential_id="gcred_abc", + organization_id="o_test", + credential_name="Default", + scopes_requested=["https://www.googleapis.com/auth/spreadsheets"], + consent_nonce="nonce-xyz", + consent_redirect_uri="http://localhost:8080/integrations/google/callback", + consent_expires_at=expires_at, + consent_code_verifier="ver-abc", + ) + + assert isinstance(result, GoogleOAuthCredentialBase) + assert result.id == "gcred_abc" + assert result.organization_id == "o_test" + assert result.credential_name == "Default" + assert result.provider == "google" + assert result.state == STATE_PENDING_CONSENT + assert result.scopes_requested == ["https://www.googleapis.com/auth/spreadsheets"] + assert result.scopes_granted == [] + assert result.created_at is not None + assert result.modified_at is not None + + +@pytest.mark.asyncio +async def test_promote_pending_to_active_returns_schema_without_greenlet_error( + repo: GoogleOAuthRepository, +) -> None: + expires_at = datetime.datetime.utcnow() + datetime.timedelta(minutes=10) + await repo.insert_pending_credential( + credential_id="gcred_promote", + organization_id="o_test", + credential_name="Default", + scopes_requested=["https://www.googleapis.com/auth/spreadsheets"], + consent_nonce="nonce-promote", + consent_redirect_uri="http://localhost:8080/integrations/google/callback", + consent_expires_at=expires_at, + consent_code_verifier="ver-promote", + ) + + result = await repo.promote_pending_to_active( + organization_id="o_test", + nonce="nonce-promote", + encrypted_refresh_token="cipher-value", + encrypted_method=EncryptMethod.AES, + scopes_granted=["https://www.googleapis.com/auth/spreadsheets"], + now=datetime.datetime.utcnow(), + ) + + assert isinstance(result, GoogleOAuthCredentialBase) + assert result.id == "gcred_promote" + assert result.state == "active" + assert result.scopes_granted == ["https://www.googleapis.com/auth/spreadsheets"] + + +@pytest.mark.asyncio +async def test_rename_active_returns_schema_without_greenlet_error( + repo: GoogleOAuthRepository, +) -> None: + expires_at = datetime.datetime.utcnow() + datetime.timedelta(minutes=10) + await repo.insert_pending_credential( + credential_id="gcred_rename", + organization_id="o_test", + credential_name="Old Name", + scopes_requested=["https://www.googleapis.com/auth/spreadsheets"], + consent_nonce="nonce-rename", + consent_redirect_uri="http://localhost:8080/integrations/google/callback", + consent_expires_at=expires_at, + consent_code_verifier="ver-rename", + ) + await repo.promote_pending_to_active( + organization_id="o_test", + nonce="nonce-rename", + encrypted_refresh_token="cipher-value", + encrypted_method=EncryptMethod.AES, + scopes_granted=["https://www.googleapis.com/auth/spreadsheets"], + now=datetime.datetime.utcnow(), + ) + + renamed = await repo.rename_active( + organization_id="o_test", + credential_id="gcred_rename", + credential_name="New Name", + now=datetime.datetime.utcnow(), + ) + + assert renamed is not None + assert isinstance(renamed, GoogleOAuthCredentialBase) + assert renamed.credential_name == "New Name" + assert renamed.state == "active" + + +@pytest.mark.asyncio +async def test_consent_app_origin_round_trips_through_load_pending_by_nonce( + repo: GoogleOAuthRepository, +) -> None: + """consent_app_origin written by insert_pending_credential is returned by load_pending_by_nonce.""" + expires_at = datetime.datetime.utcnow() + datetime.timedelta(minutes=10) + await repo.insert_pending_credential( + credential_id="gcred_app_origin", + organization_id="o_test", + credential_name="Default", + scopes_requested=["https://www.googleapis.com/auth/spreadsheets"], + consent_nonce="nonce-app-origin", + consent_redirect_uri="https://app-staging.skyvern.com/integrations/google/callback", + consent_expires_at=expires_at, + consent_code_verifier="ver-app-origin", + consent_app_origin="https://skyvern-cloud-git-branch-skyvern.vercel.app", + ) + + from skyvern.forge.sdk.db.repositories.google_oauth import PendingConsentContext + + ctx = await repo.load_pending_by_nonce(organization_id="o_test", nonce="nonce-app-origin") + assert ctx is not None + assert isinstance(ctx, PendingConsentContext) + assert ctx.consent_app_origin == "https://skyvern-cloud-git-branch-skyvern.vercel.app" + + +@pytest.mark.asyncio +async def test_consent_app_origin_defaults_to_none_for_backward_compat( + repo: GoogleOAuthRepository, +) -> None: + """Omitting consent_app_origin (pre-existing callers) stores and returns None.""" + expires_at = datetime.datetime.utcnow() + datetime.timedelta(minutes=10) + await repo.insert_pending_credential( + credential_id="gcred_no_origin", + organization_id="o_test", + credential_name="Default", + scopes_requested=["https://www.googleapis.com/auth/spreadsheets"], + consent_nonce="nonce-no-origin", + consent_redirect_uri="https://app-staging.skyvern.com/integrations/google/callback", + consent_expires_at=expires_at, + consent_code_verifier="ver-no-origin", + # consent_app_origin intentionally omitted + ) + + ctx = await repo.load_pending_by_nonce(organization_id="o_test", nonce="nonce-no-origin") + assert ctx is not None + assert ctx.consent_app_origin is None + + +@pytest.mark.asyncio +async def test_load_pending_by_nonce_filters_expired_rows( + repo: GoogleOAuthRepository, +) -> None: + """Expired consent rows must not load — otherwise the callback exchanges Google's + one-time auth code before the nonce is rejected, forcing the user to restart.""" + expired_at = datetime.datetime.utcnow() - datetime.timedelta(minutes=1) + await repo.insert_pending_credential( + credential_id="gcred_expired", + organization_id="o_test", + credential_name="Default", + scopes_requested=["https://www.googleapis.com/auth/spreadsheets"], + consent_nonce="nonce-expired", + consent_redirect_uri="https://app/callback", + consent_expires_at=expired_at, + consent_code_verifier="ver-expired", + ) + + ctx = await repo.load_pending_by_nonce(organization_id="o_test", nonce="nonce-expired") + assert ctx is None diff --git a/tests/unit/google/test_sheets.py b/tests/unit/google/test_sheets.py new file mode 100644 index 000000000..d4229603f --- /dev/null +++ b/tests/unit/google/test_sheets.py @@ -0,0 +1,318 @@ +"""Tests for skyvern.forge.sdk.services.google_sheets_service. + +Uses httpx.MockTransport to fake the Drive v3 + Sheets v4 endpoints so we can +verify parameter encoding, id extraction, response parsing, and error mapping +(including the 403 insufficient_scope -> reconnect_required translation) end +to end without hitting Google. +""" + +from typing import Any, Callable +from unittest.mock import patch + +import httpx +import pytest + +from skyvern.forge.sdk.services import google_sheets_service + + +def _install_transport( + monkeypatch: pytest.MonkeyPatch, + handler: Callable[[httpx.Request], httpx.Response], +) -> None: + transport = httpx.MockTransport(handler) + real_client = httpx.AsyncClient + + def fake_async_client(*args: Any, **kwargs: Any) -> httpx.AsyncClient: + kwargs["transport"] = transport + return real_client(*args, **kwargs) + + monkeypatch.setattr(google_sheets_service.httpx, "AsyncClient", fake_async_client) + + +def test_extract_spreadsheet_id_from_url() -> None: + url = "https://docs.google.com/spreadsheets/d/1AbC-XyZ_123/edit#gid=0" + assert google_sheets_service.extract_spreadsheet_id(url) == "1AbC-XyZ_123" + + +def test_extract_spreadsheet_id_bare() -> None: + bare = "1AbCdEfGhIjKlMnOpQrSt_12345" + assert google_sheets_service.extract_spreadsheet_id(bare) == bare + + +def test_extract_spreadsheet_id_invalid() -> None: + with pytest.raises(ValueError): + google_sheets_service.extract_spreadsheet_id("") + with pytest.raises(ValueError): + google_sheets_service.extract_spreadsheet_id("not-a-sheet") + + +def test_build_drive_q_escapes_user_input() -> None: + q = google_sheets_service._build_drive_q("Jen's sheet") + assert "Jen\\'s sheet" in q + assert "mimeType = 'application/vnd.google-apps.spreadsheet'" in q + assert "trashed = false" in q + + +@pytest.mark.asyncio +async def test_list_spreadsheets_sends_expected_params(monkeypatch: pytest.MonkeyPatch) -> None: + captured: dict[str, Any] = {} + + def handler(request: httpx.Request) -> httpx.Response: + captured["url"] = str(request.url) + captured["auth"] = request.headers.get("Authorization") + return httpx.Response( + 200, + json={ + "nextPageToken": "tok-2", + "files": [ + { + "id": "1xyz", + "name": "Leads", + "modifiedTime": "2026-04-10T00:00:00Z", + "webViewLink": "https://docs.google.com/spreadsheets/d/1xyz", + }, + {"id": "", "name": "skip-empty-id"}, + ], + }, + ) + + _install_transport(monkeypatch, handler) + + paged = await google_sheets_service.list_spreadsheets( + access_token="at-1", + query="leads", + page_token="tok-1", + page_size=50, + ) + + assert paged.next_page_token == "tok-2" + assert len(paged.spreadsheets) == 1 + assert paged.spreadsheets[0].id == "1xyz" + assert paged.spreadsheets[0].name == "Leads" + assert captured["auth"] == "Bearer at-1" + assert "pageSize=50" in captured["url"] + assert "pageToken=tok-1" in captured["url"] + # Shared-drive support is required so picker results include sheets + # owned by shared drives, not just the user's My Drive. + assert "supportsAllDrives=true" in captured["url"] + assert "includeItemsFromAllDrives=true" in captured["url"] + + +@pytest.mark.asyncio +async def test_list_spreadsheets_403_insufficient_scope_maps_to_reconnect(monkeypatch: pytest.MonkeyPatch) -> None: + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response( + 403, + json={ + "error": { + "message": "Request had insufficient authentication scopes.", + "errors": [{"reason": "insufficientPermissions"}], + } + }, + ) + + _install_transport(monkeypatch, handler) + + with pytest.raises(google_sheets_service.GoogleSheetsAPIError) as exc_info: + await google_sheets_service.list_spreadsheets(access_token="at") + + assert exc_info.value.status == 403 + assert exc_info.value.code == "reconnect_required" + + +@pytest.mark.asyncio +async def test_get_spreadsheet_tabs_parses_sheet_metadata(monkeypatch: pytest.MonkeyPatch) -> None: + def handler(request: httpx.Request) -> httpx.Response: + assert "/v4/spreadsheets/1xyz" in str(request.url) + return httpx.Response( + 200, + json={ + "sheets": [ + {"properties": {"sheetId": 0, "title": "Sheet1", "index": 0}}, + {"properties": {"sheetId": 123, "title": "Q2", "index": 1}}, + {"properties": {"title": "no-id-skip"}}, + ] + }, + ) + + _install_transport(monkeypatch, handler) + + tabs = await google_sheets_service.get_spreadsheet_tabs( + access_token="at", + spreadsheet_id="https://docs.google.com/spreadsheets/d/1xyz/edit", + ) + + assert [t.title for t in tabs] == ["Sheet1", "Q2"] + assert [t.sheet_id for t in tabs] == [0, 123] + + +@pytest.mark.asyncio +async def test_create_spreadsheet_posts_title(monkeypatch: pytest.MonkeyPatch) -> None: + captured: dict[str, Any] = {} + + def handler(request: httpx.Request) -> httpx.Response: + captured["body"] = request.content.decode() + return httpx.Response( + 200, + json={ + "spreadsheetId": "1newid", + "properties": {"title": "My Sheet"}, + "spreadsheetUrl": "https://docs.google.com/spreadsheets/d/1newid", + "sheets": [{"properties": {"sheetId": 0, "title": "Sheet1", "index": 0}}], + }, + ) + + _install_transport(monkeypatch, handler) + + created = await google_sheets_service.create_spreadsheet(access_token="at", title="My Sheet") + + assert created.id == "1newid" + assert created.title == "My Sheet" + assert created.web_view_link and "1newid" in created.web_view_link + assert created.first_sheet_name == "Sheet1" + assert '"My Sheet"' in captured["body"] + + +@pytest.mark.asyncio +async def test_create_sheet_tab_returns_new_tab(monkeypatch: pytest.MonkeyPatch) -> None: + def handler(request: httpx.Request) -> httpx.Response: + assert ":batchUpdate" in str(request.url) + return httpx.Response( + 200, + json={"replies": [{"addSheet": {"properties": {"sheetId": 55, "title": "New Tab", "index": 3}}}]}, + ) + + _install_transport(monkeypatch, handler) + + tab = await google_sheets_service.create_sheet_tab( + access_token="at", + spreadsheet_id="1AbCdEfGhIjKlMnOpQrSt_12345", + title="New Tab", + ) + + assert tab.sheet_id == 55 + assert tab.title == "New Tab" + assert tab.index == 3 + + +@pytest.mark.asyncio +async def test_api_error_maps_status_and_message(monkeypatch: pytest.MonkeyPatch) -> None: + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response( + 404, + json={"error": {"message": "Requested entity was not found.", "errors": [{"reason": "notFound"}]}}, + ) + + _install_transport(monkeypatch, handler) + + with pytest.raises(google_sheets_service.GoogleSheetsAPIError) as exc_info: + await google_sheets_service.get_spreadsheet_tabs( + access_token="at", spreadsheet_id="1AbCdEfGhIjKlMnOpQrSt_12345" + ) + + assert exc_info.value.status == 404 + assert "not found" in exc_info.value.message.lower() + + +@pytest.mark.asyncio +async def test_get_sheet_grid_properties_returns_dimensions(monkeypatch: pytest.MonkeyPatch) -> None: + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, + json={ + "sheets": [ + { + "properties": { + "sheetId": 0, + "title": "Other", + "gridProperties": {"columnCount": 50, "rowCount": 1000}, + } + }, + { + "properties": { + "sheetId": 42, + "title": "Target", + "gridProperties": {"columnCount": 26, "rowCount": 1000}, + } + }, + ] + }, + ) + + _install_transport(monkeypatch, handler) + + grid = await google_sheets_service.get_sheet_grid_properties( + access_token="at", + spreadsheet_id="1AbCdEfGhIjKlMnOpQrSt_12345", + sheet_title="Target", + ) + + assert grid is not None + assert grid.sheet_id == 42 + assert grid.title == "Target" + assert grid.column_count == 26 + assert grid.row_count == 1000 + + +@pytest.mark.asyncio +async def test_get_sheet_grid_properties_by_id_matches_numeric_id(monkeypatch: pytest.MonkeyPatch) -> None: + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, + json={ + "sheets": [ + { + "properties": { + "sheetId": 7, + "title": "Target", + "gridProperties": {"columnCount": 26, "rowCount": 1000}, + } + } + ] + }, + ) + + _install_transport(monkeypatch, handler) + + grid = await google_sheets_service.get_sheet_grid_properties_by_id( + access_token="at", + spreadsheet_id="1AbCdEfGhIjKlMnOpQrSt_12345", + sheet_id=7, + ) + + assert grid is not None + assert grid.sheet_id == 7 + assert grid.title == "Target" + assert grid.column_count == 26 + + +@pytest.mark.asyncio +async def test_get_sheet_grid_properties_returns_none_when_missing(monkeypatch: pytest.MonkeyPatch) -> None: + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response(200, json={"sheets": []}) + + _install_transport(monkeypatch, handler) + + grid = await google_sheets_service.get_sheet_grid_properties( + access_token="at", + spreadsheet_id="1AbCdEfGhIjKlMnOpQrSt_12345", + sheet_title="Missing", + ) + + assert grid is None + + +def test_build_append_dimension_request_validates_inputs() -> None: + from skyvern.schemas.google_sheets import build_append_dimension_request + + req = build_append_dimension_request(sheet_id=42, dimension="COLUMNS", length=5) + assert req == {"appendDimension": {"sheetId": 42, "dimension": "COLUMNS", "length": 5}} + + with pytest.raises(ValueError): + build_append_dimension_request(sheet_id=42, dimension="DIAGONAL", length=1) + with pytest.raises(ValueError): + build_append_dimension_request(sheet_id=42, dimension="ROWS", length=0) + + +# Ensure patch is not flagged as unused by linters in case of future use +_ = patch diff --git a/tests/unit/google/test_sheets_headers.py b/tests/unit/google/test_sheets_headers.py new file mode 100644 index 000000000..4202b329d --- /dev/null +++ b/tests/unit/google/test_sheets_headers.py @@ -0,0 +1,161 @@ +"""Tests for skyvern.forge.sdk.services.google_sheets_service.get_sheet_headers and the +GET /spreadsheets/{id}/sheets/{title}/headers route. + +Follows the house pattern: httpx.MockTransport installed via monkeypatch. +""" + +from typing import Any, Callable + +import httpx +import pytest + +from skyvern.forge.sdk.services import google_sheets_service + + +def _install_transport( + monkeypatch: pytest.MonkeyPatch, + handler: Callable[[httpx.Request], httpx.Response], +) -> None: + transport = httpx.MockTransport(handler) + real_client = httpx.AsyncClient + + def fake_async_client(*args: Any, **kwargs: Any) -> httpx.AsyncClient: + kwargs["transport"] = transport + return real_client(*args, **kwargs) + + monkeypatch.setattr(google_sheets_service.httpx, "AsyncClient", fake_async_client) + + +@pytest.mark.asyncio +async def test_get_sheet_headers_returns_row_one_keyed_by_letter( + monkeypatch: pytest.MonkeyPatch, +) -> None: + def handler(request: httpx.Request) -> httpx.Response: + assert "/v4/spreadsheets/1abcdefghijklmnopqrst_12345/values/%27Sheet1%27%21A1%3AZZZ1" in str(request.url) + return httpx.Response( + 200, + json={"range": "Sheet1!A1:ZZZ1", "values": [["Name", "Email", "Date"]]}, + ) + + _install_transport(monkeypatch, handler) + + headers = await google_sheets_service.get_sheet_headers( + access_token="tok", + spreadsheet_id="1abcdefghijklmnopqrst_12345", + sheet_title="Sheet1", + ) + + assert headers == [ + google_sheets_service.SheetHeader(letter="A", name="Name"), + google_sheets_service.SheetHeader(letter="B", name="Email"), + google_sheets_service.SheetHeader(letter="C", name="Date"), + ] + + +@pytest.mark.asyncio +async def test_get_sheet_headers_returns_empty_when_no_values( + monkeypatch: pytest.MonkeyPatch, +) -> None: + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response(200, json={"range": "Tab!A1:ZZZ1"}) + + _install_transport(monkeypatch, handler) + + headers = await google_sheets_service.get_sheet_headers( + access_token="tok", + spreadsheet_id="1abcdefghijklmnopqrst_12345", + sheet_title="Tab", + ) + + assert headers == [] + + +@pytest.mark.asyncio +async def test_get_sheet_headers_skips_blank_header_cells( + monkeypatch: pytest.MonkeyPatch, +) -> None: + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response(200, json={"values": [["Name", "", "Date"]]}) + + _install_transport(monkeypatch, handler) + + headers = await google_sheets_service.get_sheet_headers( + access_token="tok", + spreadsheet_id="1abcdefghijklmnopqrst_12345", + sheet_title="S", + ) + + assert headers == [ + google_sheets_service.SheetHeader(letter="A", name="Name"), + google_sheets_service.SheetHeader(letter="C", name="Date"), + ] + + +@pytest.mark.asyncio +async def test_get_sheet_headers_translates_403_insufficient_scope( + monkeypatch: pytest.MonkeyPatch, +) -> None: + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response( + 403, + json={ + "error": { + "message": "Request had insufficient authentication scopes.", + "errors": [{"reason": "insufficientPermissions"}], + } + }, + ) + + _install_transport(monkeypatch, handler) + + with pytest.raises(google_sheets_service.GoogleSheetsAPIError) as exc_info: + await google_sheets_service.get_sheet_headers( + access_token="tok", + spreadsheet_id="1abcdefghijklmnopqrst_12345", + sheet_title="S", + ) + assert exc_info.value.status == 403 + assert exc_info.value.code == "reconnect_required" + + +@pytest.mark.asyncio +async def test_get_sheet_headers_quotes_titles_with_spaces( + monkeypatch: pytest.MonkeyPatch, +) -> None: + captured: dict[str, str] = {} + + def handler(request: httpx.Request) -> httpx.Response: + captured["url"] = str(request.url) + return httpx.Response(200, json={"values": [["Name"]]}) + + _install_transport(monkeypatch, handler) + + await google_sheets_service.get_sheet_headers( + access_token="tok", + spreadsheet_id="1abcdefghijklmnopqrst_12345", + sheet_title="Q1 Leads", + ) + # Quoted 'Q1 Leads' URL-encodes to %27Q1%20Leads%27 (or %27Q1+Leads%27) + !A1:ZZZ1 + assert "%27Q1" in captured["url"] and "Leads%27%21A1%3AZZZ1" in captured["url"] + + +@pytest.mark.asyncio +async def test_get_sheet_headers_generates_letters_beyond_Z( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Sanity-check the column-index-to-letter conversion wraps past Z.""" + row = ["h" + str(i) for i in range(28)] # AA and AB columns + expected_letters = [chr(ord("A") + i) for i in range(26)] + ["AA", "AB"] + + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response(200, json={"values": [row]}) + + _install_transport(monkeypatch, handler) + + headers = await google_sheets_service.get_sheet_headers( + access_token="tok", + spreadsheet_id="1abcdefghijklmnopqrst_12345", + sheet_title="S", + ) + + assert [h.letter for h in headers] == expected_letters diff --git a/tests/unit/google/test_sheets_retry.py b/tests/unit/google/test_sheets_retry.py new file mode 100644 index 000000000..c87bd8e9e --- /dev/null +++ b/tests/unit/google/test_sheets_retry.py @@ -0,0 +1,228 @@ +"""Verify the Sheets client retries 429 + 5xx with exponential backoff. + +Uses httpx.MockTransport to emit a 429 with Retry-After, then a 200. The +helper must honor Retry-After, backoff, and eventually succeed. +""" + +from datetime import datetime, timezone +from typing import Any, Callable + +import httpx +import pytest + +from skyvern.forge.sdk.services import google_sheets_service + + +def _install_transport(monkeypatch: pytest.MonkeyPatch, handler: Callable[[httpx.Request], httpx.Response]) -> None: + transport = httpx.MockTransport(handler) + real_client = httpx.AsyncClient + + def fake_async_client(*args: Any, **kwargs: Any) -> httpx.AsyncClient: + kwargs["transport"] = transport + return real_client(*args, **kwargs) + + monkeypatch.setattr(google_sheets_service.httpx, "AsyncClient", fake_async_client) + + +@pytest.mark.asyncio +async def test_values_append_does_not_retry_on_429(monkeypatch: pytest.MonkeyPatch) -> None: + """POST mutations must not retry on 429 either: Google does not guarantee + the quota check rejected the request before the write landed.""" + calls = {"n": 0} + sleeps: list[float] = [] + + async def fake_sleep(delay: float) -> None: + sleeps.append(delay) + + monkeypatch.setattr(google_sheets_service.asyncio, "sleep", fake_sleep) + + def handler(request: httpx.Request) -> httpx.Response: + calls["n"] += 1 + return httpx.Response(429, headers={"Retry-After": "2"}, json={"error": {"message": "rate limit"}}) + + _install_transport(monkeypatch, handler) + + with pytest.raises(google_sheets_service.GoogleSheetsAPIError) as exc: + await google_sheets_service.values_append( + access_token="tok", + spreadsheet_id="1abc_xyz_1234567890ABCDEF", + range_="Sheet1!A1", + values=[["a"]], + ) + assert exc.value.status == 429 + assert calls["n"] == 1 + assert sleeps == [] + + +@pytest.mark.asyncio +async def test_values_get_retries_on_429(monkeypatch: pytest.MonkeyPatch) -> None: + """GET is idempotent so 429 still retries with Retry-After backoff.""" + calls = {"n": 0} + sleeps: list[float] = [] + + async def fake_sleep(delay: float) -> None: + sleeps.append(delay) + + monkeypatch.setattr(google_sheets_service.asyncio, "sleep", fake_sleep) + + def handler(request: httpx.Request) -> httpx.Response: + calls["n"] += 1 + if calls["n"] == 1: + return httpx.Response(429, headers={"Retry-After": "2"}, json={"error": {"message": "rate limit"}}) + return httpx.Response(200, json={"values": [["a"]]}) + + _install_transport(monkeypatch, handler) + + payload = await google_sheets_service.values_get( + access_token="tok", + spreadsheet_id="1abc_xyz_1234567890ABCDEF", + ranges="Sheet1!A1", + ) + assert payload == {"values": [["a"]]} + assert calls["n"] == 2 + assert sleeps == [2.0] + + +@pytest.mark.asyncio +async def test_values_get_honors_http_date_retry_after(monkeypatch: pytest.MonkeyPatch) -> None: + """RFC 7231 allows Retry-After as an HTTP-date; we must sleep until that instant + instead of falling through to exponential backoff.""" + calls = {"n": 0} + sleeps: list[float] = [] + + async def fake_sleep(delay: float) -> None: + sleeps.append(delay) + + monkeypatch.setattr(google_sheets_service.asyncio, "sleep", fake_sleep) + + frozen_now = datetime(2026, 4, 21, 12, 0, 0, tzinfo=timezone.utc) + + class _FrozenDatetime(datetime): + @classmethod + def now(cls, tz: Any = None) -> datetime: + return frozen_now if tz is None else frozen_now.astimezone(tz) + + monkeypatch.setattr(google_sheets_service, "datetime", _FrozenDatetime) + + def handler(request: httpx.Request) -> httpx.Response: + calls["n"] += 1 + if calls["n"] == 1: + return httpx.Response( + 429, + headers={"Retry-After": "Tue, 21 Apr 2026 12:00:05 GMT"}, + json={"error": {"message": "rate limit"}}, + ) + return httpx.Response(200, json={"values": [["a"]]}) + + _install_transport(monkeypatch, handler) + + payload = await google_sheets_service.values_get( + access_token="tok", + spreadsheet_id="1abc_xyz_1234567890ABCDEF", + ranges="Sheet1!A1", + ) + assert payload == {"values": [["a"]]} + assert calls["n"] == 2 + assert sleeps == [5.0] + + +@pytest.mark.asyncio +async def test_values_get_retries_on_transport_error(monkeypatch: pytest.MonkeyPatch) -> None: + """GET is idempotent so transient transport failures retry with backoff.""" + calls = {"n": 0} + sleeps: list[float] = [] + + async def fake_sleep(delay: float) -> None: + sleeps.append(delay) + + monkeypatch.setattr(google_sheets_service.asyncio, "sleep", fake_sleep) + + def handler(request: httpx.Request) -> httpx.Response: + calls["n"] += 1 + if calls["n"] == 1: + raise httpx.ConnectError("connection refused", request=request) + return httpx.Response(200, json={"sheets": []}) + + _install_transport(monkeypatch, handler) + + payload = await google_sheets_service.values_get( + access_token="tok", + spreadsheet_id="1abc_xyz_1234567890ABCDEF", + ranges="Sheet1!A1", + ) + assert payload == {"sheets": []} + assert calls["n"] == 2 + assert sleeps == [1.0] + + +@pytest.mark.asyncio +async def test_values_append_does_not_replay_post_on_transport_error(monkeypatch: pytest.MonkeyPatch) -> None: + """POST mutations must not replay on transport failures: Google may have + processed the write before the response was lost, and a retry would + duplicate rows.""" + calls = {"n": 0} + + async def fake_sleep(delay: float) -> None: + return None + + monkeypatch.setattr(google_sheets_service.asyncio, "sleep", fake_sleep) + + def handler(request: httpx.Request) -> httpx.Response: + calls["n"] += 1 + raise httpx.ReadTimeout("response lost", request=request) + + _install_transport(monkeypatch, handler) + + with pytest.raises(google_sheets_service.GoogleSheetsAPIError): + await google_sheets_service.values_append( + access_token="tok", + spreadsheet_id="1abc_xyz_1234567890ABCDEF", + range_="Sheet1!A1", + values=[["a"]], + ) + assert calls["n"] == 1 # one attempt, no retry + + +@pytest.mark.asyncio +async def test_values_append_terminal_transport_error_becomes_api_error(monkeypatch: pytest.MonkeyPatch) -> None: + async def fake_sleep(delay: float) -> None: + return None + + monkeypatch.setattr(google_sheets_service.asyncio, "sleep", fake_sleep) + + def handler(request: httpx.Request) -> httpx.Response: + raise httpx.ConnectError("connection refused", request=request) + + _install_transport(monkeypatch, handler) + + with pytest.raises(google_sheets_service.GoogleSheetsAPIError) as exc: + await google_sheets_service.values_append( + access_token="tok", + spreadsheet_id="1abc_xyz_1234567890ABCDEF", + range_="Sheet1!A1", + values=[["a"]], + ) + assert exc.value.status == 503 + assert exc.value.code == "upstream_unavailable" + + +@pytest.mark.asyncio +async def test_values_append_gives_up_after_max_attempts(monkeypatch: pytest.MonkeyPatch) -> None: + async def fake_sleep(delay: float) -> None: + return None + + monkeypatch.setattr(google_sheets_service.asyncio, "sleep", fake_sleep) + + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response(503, json={"error": {"message": "upstream down"}}) + + _install_transport(monkeypatch, handler) + + with pytest.raises(google_sheets_service.GoogleSheetsAPIError) as exc: + await google_sheets_service.values_append( + access_token="tok", + spreadsheet_id="1abc_xyz_1234567890ABCDEF", + range_="Sheet1!A1", + values=[["a"]], + ) + assert exc.value.status == 503 diff --git a/tests/unit/google/test_sheets_route_helpers.py b/tests/unit/google/test_sheets_route_helpers.py new file mode 100644 index 000000000..50f5a66f4 --- /dev/null +++ b/tests/unit/google/test_sheets_route_helpers.py @@ -0,0 +1,60 @@ +"""Tests for skyvern.forge.sdk.routes.google_sheets helper functions. + +Focuses on _mint_access_token's exception triage so DB outages don't get +mislabeled as user reconnect prompts. +""" + +from unittest.mock import AsyncMock, MagicMock + +import pytest +from fastapi import HTTPException +from sqlalchemy.exc import OperationalError + +from skyvern.forge.sdk.routes import google_sheets as sheets_routes + + +def _patch_session(monkeypatch: pytest.MonkeyPatch) -> None: + """MagicMock's __aexit__ returns truthy by default, which would suppress + in-block exceptions; pin it to None so raises propagate to the caller.""" + from skyvern.forge import app + + monkeypatch.setattr( + app, + "DATABASE", + MagicMock( + Session=lambda: MagicMock( + __aenter__=AsyncMock(return_value=MagicMock()), + __aexit__=AsyncMock(return_value=None), + ) + ), + ) + + +@pytest.mark.asyncio +async def test_mint_access_token_returns_503_on_database_error(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr( + "skyvern.forge.sdk.services.google_oauth_service.load_credential_secrets", + AsyncMock(side_effect=OperationalError("SELECT 1", {}, BaseException("connection reset"))), + ) + _patch_session(monkeypatch) + + with pytest.raises(HTTPException) as exc_info: + await sheets_routes._mint_access_token("o_1", "cred_1") + + assert exc_info.value.status_code == 503 + assert exc_info.value.detail == "Database unavailable" + + +@pytest.mark.asyncio +async def test_mint_access_token_returns_409_reconnect_on_decrypt_failure(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr( + "skyvern.forge.sdk.services.google_oauth_service.load_credential_secrets", + AsyncMock(side_effect=Exception("Failed to decrypt token")), + ) + _patch_session(monkeypatch) + + with pytest.raises(HTTPException) as exc_info: + await sheets_routes._mint_access_token("o_1", "cred_1") + + assert exc_info.value.status_code == 409 + assert exc_info.value.detail["code"] == "reconnect_required" diff --git a/tests/unit/google/test_sheets_runtime.py b/tests/unit/google/test_sheets_runtime.py new file mode 100644 index 000000000..243c316c9 --- /dev/null +++ b/tests/unit/google/test_sheets_runtime.py @@ -0,0 +1,95 @@ +"""Runtime-side Sheets verbs (values.get / append / update, batchUpdate). + +These are the methods the workflow runtime calls via app.AGENT_FUNCTION; the +picker UX already covered in test_sheets.py. +""" + +from typing import Any, Callable + +import httpx +import pytest + +from skyvern.forge.sdk.services import google_sheets_service + + +def _install_transport(monkeypatch: pytest.MonkeyPatch, handler: Callable[[httpx.Request], httpx.Response]) -> None: + transport = httpx.MockTransport(handler) + real_client = httpx.AsyncClient + + def fake_async_client(*args: Any, **kwargs: Any) -> httpx.AsyncClient: + kwargs["transport"] = transport + return real_client(*args, **kwargs) + + monkeypatch.setattr(google_sheets_service.httpx, "AsyncClient", fake_async_client) + + +@pytest.mark.asyncio +async def test_values_get_returns_payload(monkeypatch: pytest.MonkeyPatch) -> None: + def handler(request: httpx.Request) -> httpx.Response: + assert "1abc_xyz_1234567890ABCDEF" in request.url.path + return httpx.Response(200, json={"sheets": [{"properties": {"title": "t"}, "data": [{"rowData": []}]}]}) + + _install_transport(monkeypatch, handler) + payload = await google_sheets_service.values_get( + access_token="tok", + spreadsheet_id="1abc_xyz_1234567890ABCDEF", + ranges="A1:B2", + fields="sheets(properties(title))", + ) + assert payload["sheets"][0]["properties"]["title"] == "t" + + +@pytest.mark.asyncio +async def test_values_append_posts_to_append_endpoint(monkeypatch: pytest.MonkeyPatch) -> None: + captured: dict[str, Any] = {} + + def handler(request: httpx.Request) -> httpx.Response: + captured["path"] = request.url.path + captured["params"] = dict(request.url.params) + return httpx.Response(200, json={"updates": {"updatedRows": 1}}) + + _install_transport(monkeypatch, handler) + payload = await google_sheets_service.values_append( + access_token="tok", + spreadsheet_id="1abc_xyz_1234567890ABCDEF", + range_="Sheet1!A1", + values=[["a", "b"]], + ) + assert payload["updates"]["updatedRows"] == 1 + assert captured["path"].endswith(":append") + assert captured["params"]["valueInputOption"] == "USER_ENTERED" + + +@pytest.mark.asyncio +async def test_batch_update_sends_requests(monkeypatch: pytest.MonkeyPatch) -> None: + def handler(request: httpx.Request) -> httpx.Response: + assert request.url.path.endswith(":batchUpdate") + return httpx.Response(200, json={"replies": [{"addSheet": {"properties": {"sheetId": 7, "title": "new"}}}]}) + + _install_transport(monkeypatch, handler) + payload = await google_sheets_service.batch_update( + access_token="tok", + spreadsheet_id="1abc_xyz_1234567890ABCDEF", + requests=[{"addSheet": {"properties": {"title": "new"}}}], + ) + assert payload["replies"][0]["addSheet"]["properties"]["sheetId"] == 7 + + +@pytest.mark.asyncio +async def test_get_sheet_id_by_title_matches(monkeypatch: pytest.MonkeyPatch) -> None: + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, + json={ + "sheets": [ + {"properties": {"sheetId": 1, "title": "A"}}, + {"properties": {"sheetId": 42, "title": "Leads"}}, + ] + }, + ) + + _install_transport(monkeypatch, handler) + sheet_id = await google_sheets_service.get_sheet_id_by_title( + access_token="tok", spreadsheet_id="1abc_xyz_1234567890ABCDEF", sheet_title="Leads" + ) + assert sheet_id == 42