gcp_utils.py

import json
import logging
import os
from google.cloud import tasks_v2
from google.protobuf.timestamp_pb2 import Timestamp

REFRESH_VIEW_TASK_EXECUTOR_BODY = json.dumps(
    {"task": "refresh_materialized_view", "payload": {"dry_run": False}}
).encode()


def create_refresh_materialized_view_task():
    """
    Asynchronously refresh a materialized view.
    Ensures deduplication by generating a unique task name.

    Returns:
        dict: Response message and status code.
    """
    from google.protobuf import timestamp_pb2
    from datetime import datetime, timedelta

    try:
        logging.info("Creating materialized view refresh task.")
        now = datetime.now()

        # BOUNCE WINDOW: next :00 or :30
        minute = now.minute
        if minute < 30:
            bucket_time = now.replace(minute=30, second=0, microsecond=0)
        else:
            bucket_time = now.replace(minute=0, second=0, microsecond=0) + timedelta(hours=1)

        timestamp_str = bucket_time.strftime("%Y-%m-%d-%H-%M")
        task_name = f"refresh-materialized-view-{timestamp_str}"

        # Convert to protobuf timestamp
        proto_time = timestamp_pb2.Timestamp()
        proto_time.FromDatetime(bucket_time)

        # Cloud Tasks setup
        project = os.getenv("PROJECT_ID")
        queue = os.getenv("MATERIALIZED_VIEW_QUEUE")
        logging.debug("Queue name from env: %s", queue)
        gcp_region = os.getenv("GCP_REGION")
        environment_name = os.getenv("ENVIRONMENT")
        url = f"https://{gcp_region}-" f"{project}.cloudfunctions.net/" f"tasks_executor-{environment_name}"
        # Enqueue the task
        try:
            create_http_task_with_name(
                client=tasks_v2.CloudTasksClient(),
                body=REFRESH_VIEW_TASK_EXECUTOR_BODY,
                url=url,
                project_id=project,
                gcp_region=gcp_region,
                queue_name=queue,
                task_name=task_name,
                task_time=proto_time,
                http_method=tasks_v2.HttpMethod.POST,
            )
            logging.info("Scheduled refresh materialized view task for %s", task_name)
            return {"message": "Refresh task for %s scheduled." % task_name}, 200
        except Exception as e:
            if "ALREADY_EXISTS" in str(e):
                logging.info("Task already exists for %s, skipping.", task_name)

    except Exception as error:
        logging.error("Error enqueuing task: %s", error)
        return {"error": "Error enqueuing task: %s" % error}, 500


def create_http_task_with_name(
    client: "tasks_v2.CloudTasksClient",
    body: bytes,
    url: str,
    project_id: str,
    gcp_region: str,
    queue_name: str,
    task_name: str,
    task_time: Timestamp,
    http_method: "tasks_v2.HttpMethod",
):
    """Creates a GCP Cloud Task."""
    token = tasks_v2.OidcToken(service_account_email=os.getenv("SERVICE_ACCOUNT_EMAIL"))

    parent = client.queue_path(project_id, gcp_region, queue_name)
    logging.info("Queue parent path: %s", parent)

    task = tasks_v2.Task(
        # If task_name is provided, it will be used; otherwise, a unique name will be generated.
        # This is useful for deduplication purposes.
        name=f"{parent}/tasks/{task_name}" if task_name else None,
        schedule_time=task_time,
        http_request=tasks_v2.HttpRequest(
            url=url,
            http_method=http_method,
            oidc_token=token,
            body=body,
            headers={"Content-Type": "application/json"},
        ),
    )
    try:
        response = client.create_task(parent=parent, task=task)
        logging.info("Task created with task_name: %s", task_name)
    except Exception as e:
        if "Requested entity already exists" in str(e):
            logging.info("Task already exists for %s, skipping.", task_name)
        else:
            logging.error("Error creating task: %s", e)
            logging.error("response: %s", response)