Repository: PrefectHQ/prefect
File: src/integrations/prefect-gcp/prefect_gcp/workers/cloud_run_v2.py (Python)

Code before:
from __future__ import annotations

import re
import shlex
import time
from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional
from uuid import uuid4

from anyio.abc import TaskStatus
from google.api_core.client_options import ClientOptions
from googleapiclient import discovery

# noinspection PyProtectedMember
from googleapiclient.discovery import Resource
from googleapiclient.errors import HttpError
from jsonpatch import JsonPatch
from pydantic import Field, PrivateAttr, field_validator

from prefect._internal.retries import exponential_backoff_with_jitter
from prefect.exceptions import InfrastructureNotFound
from prefect.logging.loggers import PrefectLogAdapter, flow_run_logger
from prefect.utilities.asyncutils import run_sync_in_worker_thread
from prefect.utilities.dockerutils import get_prefect_image_name
from prefect.workers.base import (
    BaseJobConfiguration,
    BaseVariables,
    BaseWorker,
    BaseWorkerResult,
)
from prefect_gcp.credentials import GcpCredentials
from prefect_gcp.models.cloud_run_v2 import ExecutionV2, JobV2, SecretKeySelector
from prefect_gcp.utilities import slugify_name

if TYPE_CHECKING:
    from prefect.client.schemas.objects import Flow, FlowRun, WorkPool
    from prefect.client.schemas.responses import DeploymentResponse


def _get_default_job_body_template() -> Dict[str, Any]:
    """
    Returns the default job body template for the Cloud Run worker.

    Returns:
        The default job body template.

Review comment (refactor):
I'd prefer not to use an internal utility from the core library in an integration, since it's intended only for the core library and could lead to future breakage.

Instead, could you use `tenacity` for retries? We already use it in the Vertex worker for a similar purpose.

Code after:
from __future__ import annotations

import re
import shlex
import time
from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional
from uuid import uuid4

from anyio.abc import TaskStatus
from google.api_core.client_options import ClientOptions
from googleapiclient import discovery

# noinspection PyProtectedMember
from googleapiclient.discovery import Resource
from googleapiclient.errors import HttpError
from jsonpatch import JsonPatch
from pydantic import Field, PrivateAttr, field_validator
from tenacity import (
    Retrying,
    retry_if_exception,
    stop_after_attempt,
    wait_exponential_jitter,
)

from prefect.exceptions import InfrastructureNotFound
from prefect.logging.loggers import PrefectLogAdapter, flow_run_logger
from prefect.utilities.asyncutils import run_sync_in_worker_thread
from prefect.utilities.dockerutils import get_prefect_image_name
from prefect.workers.base import (
    BaseJobConfiguration,
    BaseVariables,
    BaseWorker,
    BaseWorkerResult,
)
from prefect_gcp.credentials import GcpCredentials
from prefect_gcp.models.cloud_run_v2 import ExecutionV2, JobV2, SecretKeySelector
from prefect_gcp.utilities import slugify_name

if TYPE_CHECKING:
    from prefect.client.schemas.objects import Flow, FlowRun, WorkPool
    from prefect.client.schemas.responses import DeploymentResponse


def _get_default_job_body_template() -> Dict[str, Any]: